diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 2dcb08b3b..56cbcd9a2 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -8,7 +8,7 @@ jobs:
     strategy:
       matrix:
         operating-system: [ubuntu-latest, windows-latest, macos-latest]
-        python-version: [3.6, 3.7, 3.8, 3.9]
+        python-version: ['3.6', '3.7', '3.8', '3.9', '3.10']
         pytorch-toolbelt-version: [tests]
       fail-fast: false
     steps:
@@ -20,16 +20,8 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Update pip
       run: python -m pip install --upgrade pip
-    - name: Install PyTorch on Linux and Windows
-      if: >
-        matrix.operating-system == 'ubuntu-latest' ||
-        matrix.operating-system == 'windows-latest'
-      run: >
-        pip install torch==1.10.1+cpu torchvision==0.11.2+cpu
-        -f https://download.pytorch.org/whl/torch_stable.html
-    - name: Install PyTorch on MacOS
-      if: matrix.operating-system == 'macos-latest'
-      run: pip install torch==1.10.1 torchvision==0.11.2
+    - name: Install PyTorch
+      run: pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Install dependencies
       run: pip install .[${{ matrix.pytorch-toolbelt-version }}]
     - name: Install linters
@@ -48,7 +40,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: ['3.6', '3.7', '3.8', '3.9', '3.10']
     steps:
     - name: Checkout
       uses: actions/checkout@v2
diff --git a/pytorch_toolbelt/__init__.py b/pytorch_toolbelt/__init__.py
index cfcbe1567..08c7a5c7f 100644
--- a/pytorch_toolbelt/__init__.py
+++ b/pytorch_toolbelt/__init__.py
@@ -1,3 +1,3 @@
 from __future__ import absolute_import
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"
diff --git a/pytorch_toolbelt/datasets/mean_std.py b/pytorch_toolbelt/datasets/mean_std.py
index 0083c1eb6..d89cfea2c 100644
--- a/pytorch_toolbelt/datasets/mean_std.py
+++ b/pytorch_toolbelt/datasets/mean_std.py
@@ -51,8 +51,8 @@ def accumulate(self, image: np.ndarray, mask: Optional[np.ndarray] = None) -> No
         image = image.reshape((-1, self.num_channels))
 
         if mask is not None:
-            mask = mask.reshape((mask.shape[0] * mask.shape[1], 1))
-            image = image[mask]
+            mask = mask.reshape((mask.shape[0] * mask.shape[1]))
+            image = image[mask, :]
 
             # In case the whole image is masked out, we exclude it entirely
             if len(image) == 0:
diff --git a/pytorch_toolbelt/utils/torch_utils.py b/pytorch_toolbelt/utils/torch_utils.py
index ff6c2adb7..eaaf15011 100644
--- a/pytorch_toolbelt/utils/torch_utils.py
+++ b/pytorch_toolbelt/utils/torch_utils.py
@@ -311,11 +311,11 @@ def describe_outputs(outputs: Union[Tensor, Dict[str, Tensor], Iterable[Tensor]]
     """
     if torch.is_tensor(outputs):
         desc = dict(size=tuple(outputs.size()), mean=outputs.mean().item(), std=outputs.std().item())
-    elif isinstance(outputs, collections.Mapping):
+    elif isinstance(outputs, collections.abc.Mapping):
         desc = {}
         for key, value in outputs.items():
             desc[key] = describe_outputs(value)
-    elif isinstance(outputs, collections.Iterable):
+    elif isinstance(outputs, collections.abc.Iterable):
         desc = []
         for index, output in enumerate(outputs):
             desc.append(describe_outputs(output))
diff --git a/setup.py b/setup.py
index 33bca8473..51a790139 100644
--- a/setup.py
+++ b/setup.py
@@ -5,6 +5,7 @@
 import sys
 
 from setuptools import find_packages, setup
+from distutils.version import LooseVersion
 
 
 def is_docker() -> bool:
@@ -88,9 +89,9 @@ def get_opencv_requirement():
 
 DEPENDENCIES = [
     # We rely on particular activation functions that were added in 1.8.1
-    "torch>=1.8.1",
+    "torch>=1.11.0" if LooseVersion(sys.version) >= LooseVersion("3.7") else "torch>=1.10.0",
     # We use some pretrained models from torchvision
-    "torchvision>=0.9.1",
+    "torchvision",
     # Library uses scipy for linear_sum_assignment for match_bboxes.
     # 1.4.0 is the first release where `maximize` argument gets introduced to this function
     "scipy>=1.4.0",
@@ -116,9 +117,7 @@ def load_readme():
 
 
 def get_test_requirements():
-    requirements = ["pytest", "onnx==1.9.0", "black==22.8.0", "timm==0.6.7", "matplotlib"]
-    if sys.version_info < (3, 3):
-        requirements.append("mock")
+    requirements = ["pytest", "black==22.8.0", "timm==0.6.7", "matplotlib"]
     return requirements
 
 
@@ -147,6 +146,10 @@ def get_test_requirements():
         "Unet",
         "Focal",
         "FPN",
+        "EfficientNet",
+        "Test-Time Augmentation",
+        "Model Ensembling",
+        "Model Surgery",
     ],
     scripts=[],
     license="License :: OSI Approved :: MIT License",
@@ -158,9 +161,14 @@ def get_test_requirements():
         "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
         "Operating System :: OS Independent",
         "Topic :: Scientific/Engineering :: Mathematics",
         "Topic :: Scientific/Engineering :: Image Recognition",
+        "Topic :: Scientific/Engineering :: Deep Learning",
+        "Topic :: Scientific/Engineering :: Computer Vision",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
         "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
diff --git a/tests/test_encoders.py b/tests/test_encoders.py
index 8c12acba2..9f681ab05 100644
--- a/tests/test_encoders.py
+++ b/tests/test_encoders.py
@@ -3,7 +3,7 @@
 
 import pytorch_toolbelt.modules.encoders as E
 from pytorch_toolbelt.modules.backbone.inceptionv4 import inceptionv4
-from pytorch_toolbelt.utils.torch_utils import maybe_cuda, count_parameters
+from pytorch_toolbelt.utils.torch_utils import maybe_cuda, count_parameters, describe_outputs
 from pytorch_toolbelt.modules.encoders import timm
 
 skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
@@ -249,6 +249,8 @@ def test_timm_encoders(encoder, encoder_params):
         assert feature_map.size(2) * expected_stride == x.size(2)
         assert feature_map.size(3) * expected_stride == x.size(3)
 
+    print(describe_outputs(output))
+
 
 @pytest.mark.parametrize(
     ["encoder", "encoder_params"],
diff --git a/tests/test_onnx.py b/tests/test_model_export.py
similarity index 92%
rename from tests/test_onnx.py
rename to tests/test_model_export.py
index 8dfac8b54..a35cc9b03 100644
--- a/tests/test_onnx.py
+++ b/tests/test_model_export.py
@@ -6,7 +6,18 @@
 from pytorch_toolbelt.utils.torch_utils import maybe_cuda, count_parameters
 from pytorch_toolbelt.modules.encoders import timm
 
-skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="Cuda is not available")
+
+def is_onnx_available():
+    try:
+        import onnx
+
+        return True
+    except ImportError:
+        return False
+
+
+skip_if_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+skip_if_no_onnx = pytest.mark.skipif(not is_onnx_available, reason="ONNX is not available")
 
 
 @pytest.mark.parametrize(
@@ -26,6 +37,7 @@
     ],
 )
 @skip_if_no_cuda
+@skip_if_no_onnx
 def test_onnx_export(encoder, encoder_params):
     import onnx