diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 645e497..46e34c9 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -12,12 +12,12 @@ jobs:
 
     steps:
       - name: Check out Git repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
 
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
-          python-version: 3.8
+          python-version: 3.9
 
       - uses: liskin/gh-problem-matcher-wrap@v1
         with:
@@ -40,29 +40,17 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
 
-      - name: Cache Conda
-        uses: actions/cache@v2
-        with:
-          path: /usr/share/miniconda/envs/torcharc
-          key: ${{ runner.os }}-conda-${{ hashFiles('environment.yml') }}
-          restore-keys: |
-            ${{ runner.os }}-conda-
-
-      - name: Setup Conda dependencies
-        uses: conda-incubator/setup-miniconda@v2
+      - name: Set up Python
+        uses: actions/setup-python@v4
         with:
-          activate-environment: torcharc
-          environment-file: environment.yml
-          python-version: 3.8
-          auto-activate-base: false
+          python-version: 3.9
 
-      - name: Conda info
-        shell: bash -l {0}
+      - name: Install repo
         run: |
-          conda info
-          conda list
+          pip install torch
+          pip install -e .
 
       - uses: liskin/gh-problem-matcher-wrap@v1
         with:
@@ -70,8 +58,9 @@ jobs:
           linters: pytest
 
       - name: Run tests
-        shell: bash -l {0}
-        run: python setup.py test | tee pytest-coverage.txt
+        run: |
+          pip install wheel
+          python setup.py test | tee pytest-coverage.txt
 
       - name: Post coverage to PR comment
         uses: coroo/pytest-coverage-commentator@v1.0.2
diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml
new file mode 100644
index 0000000..dec4f78
--- /dev/null
+++ b/.github/workflows/publish.yaml
@@ -0,0 +1,31 @@
+name: Publish to PyPI
+
+on:
+  push:
+    branches: [main]
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.9
+
+      - name: Install repo
+        run: pip install -e .
+
+      - name: Build a binary wheel
+        run: |
+          pip install wheel
+          python setup.py bdist_wheel
+
+      - name: Publish package
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_API_TOKEN }}
diff --git a/.github/workflows/tag-release.yml b/.github/workflows/tag-release.yml
index b187658..cd5a005 100644
--- a/.github/workflows/tag-release.yml
+++ b/.github/workflows/tag-release.yml
@@ -11,14 +11,14 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
 
       - name: Bump version and push tag
         id: tag_version
         uses: mathieudutour/github-tag-action@v5.1
         with:
           github_token: ${{ secrets.GITHUB_TOKEN }}
-          tag_prefix: ''
+          tag_prefix: ""
 
       - name: Create a GitHub release
         uses: actions/create-release@v1
diff --git a/bin/setup b/bin/setup
index 333569c..90cbd9f 100755
--- a/bin/setup
+++ b/bin/setup
@@ -25,7 +25,7 @@ echo "--- Installing Conda environment ---"
 if conda env list | grep "^torcharc " >/dev/null; then
     echo "conda env torcharc is already installed"
 else
-    conda create -n torcharc python=3.8.2 -y
+    conda create -n torcharc python=3.9.0 -y
     # Install ipykernel
     conda activate torcharc
     conda install ipykernel -y
diff --git a/environment.yml b/environment.yml
deleted file mode 100644
index 45b1391..0000000
--- a/environment.yml
+++ /dev/null
@@ -1,15 +0,0 @@
-name: torcharc
-channels:
-  - pytorch
-  - conda-forge
-  - defaults
-dependencies:
-  - flaky=3.6.1
-  - numpy=1.21.2
-  - pytest=5.4.1
-  - pytest-cov=2.8.1
-  - pytest-timeout=1.3.4
-  - pytorch=1.9.1
-  - pip:
-      - einops==0.3.2
-      - pydash==4.7.6
diff --git a/setup.py b/setup.py
index 8e20905..19ec18f 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,6 @@
     '--log-level=INFO',
     '--log-cli-level=INFO',
     '--log-file-level=INFO',
-    '--no-flaky-report',
     '--timeout=300',
     '--cov=torcharc',
     'test',
@@ -32,7 +31,7 @@ def run_tests(self):
 
 setup(
     name='torcharc',
-    version='1.2.0',
+    version='1.2.1',
     description='Build PyTorch networks by specifying architectures.',
     long_description='https://github.com/kengz/torcharc',
     keywords='torcharc',
@@ -53,8 +52,7 @@ def run_tests(self):
     tests_require=[
         'autopep8==1.5.3',
         'flake8==3.8.3',
-        'flaky==3.6.1',
-        'pytest==5.4.1',
+        'pytest',
         'pytest-cov==2.8.1',
         'pytest-timeout==1.3.4',
     ],
diff --git a/test/module/perceiver_io/test_preprocessor.py b/test/module/perceiver_io/test_preprocessor.py
index 1f176bf..12a3f8b 100644
--- a/test/module/perceiver_io/test_preprocessor.py
+++ b/test/module/perceiver_io/test_preprocessor.py
@@ -12,7 +12,6 @@
 def test_text_preprocessor(batch, vocab_size, embed_dim, max_seq_len):
     x = torch.randint(vocab_size, (batch, max_seq_len))
     module = preprocessor.TextPreprocessor(vocab_size, embed_dim, max_seq_len)
-    assert embed_dim == module.output_dim
     out = module(x)
     assert [max_seq_len, embed_dim] == module.out_shape
     assert list(out.shape) == [batch, *module.out_shape]
@@ -26,7 +25,7 @@ def test_text_preprocessor(batch, vocab_size, embed_dim, max_seq_len):
     [64, 64, 64, 3]  # volume
 ])
 @pytest.mark.parametrize('num_freq_bands', [32])
-def test_text_preprocessor(batch, in_shape, num_freq_bands):
+def test_fourier_preprocessor(batch, in_shape, num_freq_bands):
     x = torch.rand(batch, *in_shape)
     max_reso = [2 * r for r in in_shape[:-1]]  # max resolution twice the input size
     module = preprocessor.FourierPreprocessor(in_shape, num_freq_bands, max_reso)
diff --git a/test/module/test_merge.py b/test/module/test_merge.py
index 767543b..ab93ddd 100644
--- a/test/module/test_merge.py
+++ b/test/module/test_merge.py
@@ -1,4 +1,3 @@
-from torcharc import net_util
 from torcharc.module.merge import ConcatMerge, FiLMMerge, Merge
 import pytest
 import torch
diff --git a/torcharc/module/perceiver_io/perceiver.py b/torcharc/module/perceiver_io/perceiver.py
index 1a64ef8..6b991ff 100644
--- a/torcharc/module/perceiver_io/perceiver.py
+++ b/torcharc/module/perceiver_io/perceiver.py
@@ -4,7 +4,6 @@
 from torcharc.module.perceiver_io import encoder, decoder, postprocessor, preprocessor
 from torcharc.module.perceiver_io.attention import SpreadSequential
 from torcharc.net_util import build_component
-from typing import Union
 import pydash as ps
 
 
diff --git a/torcharc/module/perceiver_io/preprocessor.py b/torcharc/module/perceiver_io/preprocessor.py
index b61ebeb..e49d111 100644
--- a/torcharc/module/perceiver_io/preprocessor.py
+++ b/torcharc/module/perceiver_io/preprocessor.py
@@ -75,7 +75,7 @@ def __init__(self, in_shape: list, num_freq_bands: int, max_reso: list = None, c
     def build_positions(self, start: float = -1.0, end: float = 1.0):
         '''Build spatial coordinates as a meshgrid, i.e. coordinates laid out such that values along the channel is a point in coordinate, e.g. shape = (x, y, 2)'''
         x_y = [torch.linspace(start, end, steps=s) for s in self.spatial_shape]
-        return torch.stack(torch.meshgrid(*x_y), dim=len(self.spatial_shape))
+        return torch.stack(torch.meshgrid(*x_y, indexing='xy'), dim=len(self.spatial_shape))
 
     def build_pos_encoding(self, pos: torch.Tensor, max_reso: list = None) -> torch.Tensor:
         '''