braindecode 0.6__tar.gz → 0.7__tar.gz

Braindecode-0.7/Braindecode.egg-info/PKG-INFO

@@ -0,0 +1,180 @@
+Metadata-Version: 2.1
+Name: Braindecode
+Version: 0.7
+Summary: A deep learning toolbox to decode raw time-domain EEG.
+Home-page: https://github.com/braindecode/braindecode
+Author: Robin Tibor Schirrmeister
+Author-email: robintibor@gmail.com
+License: BSD 3-Clause
+Keywords: eeg deep-learning brain-state-decoding
+Platform: UNKNOWN
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Software Development :: Build Tools
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3.7
+License-File: LICENSE.txt
+
+.. image:: https://badges.gitter.im/braindecodechat/community.svg
+   :alt: Join the chat at https://gitter.im/braindecodechat/community
+   :target: https://gitter.im/braindecodechat/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge
+
+.. image:: https://github.com/braindecode/braindecode/workflows/tests-and-docs/badge.svg
+   :target: https://github.com/braindecode/braindecode/actions
+
+.. image:: https://circleci.com/gh/braindecode/braindecode.svg?style=svg
+   :target: https://circleci.com/gh/braindecode/braindecode
+   :alt: Doc build on CircleCI
+
+.. image:: https://codecov.io/gh/braindecode/braindecode/branch/master/graph/badge.svg
+   :target: https://codecov.io/gh/braindecode/braindecode
+   :alt: Code Coverage
+
+.. |Braindecode| image:: https://user-images.githubusercontent.com/42702466/177958779-b00628aa-9155-4c51-96d1-d8c345aff575.svg
+.. _braindecode: braindecode.org/
+
+Braindecode
+===========
+
+Braindecode is an open-source Python toolbox for decoding raw electrophysiological brain
+data with deep learning models. It includes dataset fetchers, data preprocessing and
+visualization tools, as well as implementations of several deep learning
+architectures and data augmentations for analysis of EEG, ECoG and MEG.
+
+For neuroscientists who want to work with deep learning and
+deep learning researchers who want to work with neurophysiological data.
+
+
+Installation Braindecode
+========================
+
+1. Install pytorch from http://pytorch.org/ (you don't need to install torchvision).
+
+2. Install `MOABB <https://github.com/NeuroTechX/moabb>`_ via pip (needed if you want to use MOABB datasets utilities):
+
+.. code-block:: bash
+
+  pip install moabb
+
+3. Install latest release of braindecode via pip:
+
+.. code-block:: bash
+
+  pip install braindecode
+
+alternatively, if you use conda, you could create a dedicated environment with the following:
+
+.. code-block:: bash
+
+  curl -O https://raw.githubusercontent.com/braindecode/braindecode/master/environment.yml
+  conda env create -f environment.yml
+  conda activate braindecode
+
+alternatively, install the latest version of braindecode via pip:
+
+.. code-block:: bash
+
+  pip install -U https://api.github.com/repos/braindecode/braindecode/zipball/master
+
+
+Documentation
+=============
+
+Documentation is online under https://braindecode.org, both in the stable and dev versions.
+
+
+Contributing to Braindecode
+===========================
+
+Guidelines for contributing to the library can be found on the braindecode github:
+
+https://github.com/braindecode/braindecode/blob/master/CONTRIBUTING.md
+
+Braindecode chat
+================
+
+https://gitter.im/braindecodechat/community
+
+Citing
+======
+
+If you use this code in a scientific publication, please cite us as:
+
+.. code-block:: bibtex
+
+  @article {HBM:HBM23730,
+  author = {Schirrmeister, Robin Tibor and Springenberg, Jost Tobias and Fiederer,
+    Lukas Dominique Josef and Glasstetter, Martin and Eggensperger, Katharina and Tangermann, Michael and
+    Hutter, Frank and Burgard, Wolfram and Ball, Tonio},
+  title = {Deep learning with convolutional neural networks for EEG decoding and visualization},
+  journal = {Human Brain Mapping},
+  issn = {1097-0193},
+  url = {http://dx.doi.org/10.1002/hbm.23730},
+  doi = {10.1002/hbm.23730},
+  month = {aug},
+  year = {2017},
+  keywords = {electroencephalography, EEG analysis, machine learning, end-to-end learning, brain–machine interface,
+    brain–computer interface, model interpretability, brain mapping},
+  }
+
+as well as the `MNE-Python <https://mne.tools>`_ software that is used by braindecode:
+
+.. code-block:: bibtex
+
+  @article{10.3389/fnins.2013.00267,
+  author={Gramfort, Alexandre and Luessi, Martin and Larson, Eric and Engemann, Denis and Strohmeier, Daniel and Brodbeck, Christian and Goj, Roman and Jas, Mainak and Brooks, Teon and Parkkonen, Lauri and Hämäläinen, Matti},
+  title={{MEG and EEG data analysis with MNE-Python}},
+  journal={Frontiers in Neuroscience},
+  volume={7},
+  pages={267},
+  year={2013},
+  url={https://www.frontiersin.org/article/10.3389/fnins.2013.00267},
+  doi={10.3389/fnins.2013.00267},
+  issn={1662-453X},
+  }
+
+
+
+
+Licensing
+^^^^^^^^^
+
+Braindecode is **BSD-licenced** (BSD-3-Clause):
+
+    This software is OSI Certified Open Source Software.
+    OSI Certified is a certification mark of the Open Source Initiative.
+
+    Copyright (c) 2011-2022, authors of Braindecode.
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions are met:
+
+    * Redistributions of source code must retain the above copyright notice,
+      this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above copyright notice,
+      this list of conditions and the following disclaimer in the documentation
+      and/or other materials provided with the distribution.
+
+    * Neither the names of braindecode authors nor the names of any
+      contributors may be used to endorse or promote products derived from
+      this software without specific prior written permission.
+
+    **This software is provided by the copyright holders and contributors
+    "as is" and any express or implied warranties, including, but not
+    limited to, the implied warranties of merchantability and fitness for
+    a particular purpose are disclaimed. In no event shall the copyright
+    owner or contributors be liable for any direct, indirect, incidental,
+    special, exemplary, or consequential damages (including, but not
+    limited to, procurement of substitute goods or services; loss of use,
+    data, or profits; or business interruption) however caused and on any
+    theory of liability, whether in contract, strict liability, or tort
+    (including negligence or otherwise) arising in any way out of the use
+    of this software, even if advised of the possibility of such
+    damage.**
+
+
+

Braindecode-0.7/Braindecode.egg-info/SOURCES.txt

@@ -0,0 +1,113 @@
+LICENSE.txt
+MANIFEST.in
+README.rst
+setup.cfg
+setup.py
+Braindecode.egg-info/PKG-INFO
+Braindecode.egg-info/SOURCES.txt
+Braindecode.egg-info/dependency_links.txt
+Braindecode.egg-info/not-zip-safe
+Braindecode.egg-info/requires.txt
+Braindecode.egg-info/top_level.txt
+braindecode/__init__.py
+braindecode/classifier.py
+braindecode/regressor.py
+braindecode/util.py
+braindecode/version.py
+braindecode/augmentation/__init__.py
+braindecode/augmentation/base.py
+braindecode/augmentation/functional.py
+braindecode/augmentation/transforms.py
+braindecode/datasets/__init__.py
+braindecode/datasets/base.py
+braindecode/datasets/bbci.py
+braindecode/datasets/bcicomp.py
+braindecode/datasets/mne.py
+braindecode/datasets/moabb.py
+braindecode/datasets/sleep_physionet.py
+braindecode/datasets/tuh.py
+braindecode/datasets/xy.py
+braindecode/datautil/__init__.py
+braindecode/datautil/mne.py
+braindecode/datautil/preprocess.py
+braindecode/datautil/serialization.py
+braindecode/datautil/util.py
+braindecode/datautil/windowers.py
+braindecode/datautil/xy.py
+braindecode/models/__init__.py
+braindecode/models/deep4.py
+braindecode/models/eeginception.py
+braindecode/models/eegitnet.py
+braindecode/models/eegnet.py
+braindecode/models/eegresnet.py
+braindecode/models/functions.py
+braindecode/models/hybrid.py
+braindecode/models/modules.py
+braindecode/models/shallow_fbcsp.py
+braindecode/models/sleep_stager_blanco_2020.py
+braindecode/models/sleep_stager_chambon_2018.py
+braindecode/models/sleep_stager_eldele_2021.py
+braindecode/models/tcn.py
+braindecode/models/tidnet.py
+braindecode/models/usleep.py
+braindecode/models/util.py
+braindecode/preprocessing/__init__.py
+braindecode/preprocessing/preprocess.py
+braindecode/preprocessing/windowers.py
+braindecode/samplers/__init__.py
+braindecode/samplers/base.py
+braindecode/samplers/ssl.py
+braindecode/training/__init__.py
+braindecode/training/callbacks.py
+braindecode/training/losses.py
+braindecode/training/scoring.py
+braindecode/visualization/__init__.py
+braindecode/visualization/confusion_matrices.py
+braindecode/visualization/gradients.py
+docs/Makefile
+docs/api.rst
+docs/cite.rst
+docs/conf.py
+docs/help.rst
+docs/index.rst
+docs/install.rst
+docs/starting.rst
+docs/whats_new.rst
+docs/_templates/autosummary/class.rst
+docs/_templates/autosummary/function.rst
+test/__init__.py
+test/acceptance_tests/__init__.py
+test/acceptance_tests/test_cropped_decoding.py
+test/acceptance_tests/test_eeg_classifier.py
+test/acceptance_tests/test_model_architectures.py
+test/acceptance_tests/test_trialwise_decoding.py
+test/acceptance_tests/test_variable_length_trials_decoding.py
+test/unit_tests/__init__.py
+test/unit_tests/test_eeg_classifier.py
+test/unit_tests/test_eeg_regressor.py
+test/unit_tests/test_util.py
+test/unit_tests/augmentation/__init__.py
+test/unit_tests/augmentation/conftest.py
+test/unit_tests/augmentation/test_base.py
+test/unit_tests/augmentation/test_transforms.py
+test/unit_tests/datasets/__init__.py
+test/unit_tests/datasets/test_bcicomp.py
+test/unit_tests/datasets/test_dataset.py
+test/unit_tests/datasets/test_mne.py
+test/unit_tests/datasets/test_sleep_physionet.py
+test/unit_tests/datasets/test_tuh.py
+test/unit_tests/datasets/test_xy.py
+test/unit_tests/datautil/__init__.py
+test/unit_tests/datautil/test_serialization.py
+test/unit_tests/models/__init__.py
+test/unit_tests/models/test_models.py
+test/unit_tests/models/test_modules.py
+test/unit_tests/models/test_util.py
+test/unit_tests/preprocessing/__init__.py
+test/unit_tests/preprocessing/test_preprocess.py
+test/unit_tests/preprocessing/test_windowers.py
+test/unit_tests/training/__init__.py
+test/unit_tests/training/test_losses.py
+test/unit_tests/training/test_scoring.py
+test/unit_tests/visualization/__init__.py
+test/unit_tests/visualization/test_gradients.py

Braindecode-0.7/PKG-INFO

@@ -0,0 +1,180 @@
+Metadata-Version: 2.1
+Name: Braindecode
+Version: 0.7
+Summary: A deep learning toolbox to decode raw time-domain EEG.
+Home-page: https://github.com/braindecode/braindecode
+Author: Robin Tibor Schirrmeister
+Author-email: robintibor@gmail.com
+License: BSD 3-Clause
+Keywords: eeg deep-learning brain-state-decoding
+Platform: UNKNOWN
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Software Development :: Build Tools
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Programming Language :: Python :: 3.7
+License-File: LICENSE.txt
+
+.. image:: https://badges.gitter.im/braindecodechat/community.svg
+   :alt: Join the chat at https://gitter.im/braindecodechat/community
+   :target: https://gitter.im/braindecodechat/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge
+
+.. image:: https://github.com/braindecode/braindecode/workflows/tests-and-docs/badge.svg
+   :target: https://github.com/braindecode/braindecode/actions
+
+.. image:: https://circleci.com/gh/braindecode/braindecode.svg?style=svg
+   :target: https://circleci.com/gh/braindecode/braindecode
+   :alt: Doc build on CircleCI
+
+.. image:: https://codecov.io/gh/braindecode/braindecode/branch/master/graph/badge.svg
+   :target: https://codecov.io/gh/braindecode/braindecode
+   :alt: Code Coverage
+
+.. |Braindecode| image:: https://user-images.githubusercontent.com/42702466/177958779-b00628aa-9155-4c51-96d1-d8c345aff575.svg
+.. _braindecode: braindecode.org/
+
+Braindecode
+===========
+
+Braindecode is an open-source Python toolbox for decoding raw electrophysiological brain
+data with deep learning models. It includes dataset fetchers, data preprocessing and
+visualization tools, as well as implementations of several deep learning
+architectures and data augmentations for analysis of EEG, ECoG and MEG.
+
+For neuroscientists who want to work with deep learning and
+deep learning researchers who want to work with neurophysiological data.
+
+
+Installation Braindecode
+========================
+
+1. Install pytorch from http://pytorch.org/ (you don't need to install torchvision).
+
+2. Install `MOABB <https://github.com/NeuroTechX/moabb>`_ via pip (needed if you want to use MOABB datasets utilities):
+
+.. code-block:: bash
+
+  pip install moabb
+
+3. Install latest release of braindecode via pip:
+
+.. code-block:: bash
+
+  pip install braindecode
+
+alternatively, if you use conda, you could create a dedicated environment with the following:
+
+.. code-block:: bash
+
+  curl -O https://raw.githubusercontent.com/braindecode/braindecode/master/environment.yml
+  conda env create -f environment.yml
+  conda activate braindecode
+
+alternatively, install the latest version of braindecode via pip:
+
+.. code-block:: bash
+
+  pip install -U https://api.github.com/repos/braindecode/braindecode/zipball/master
+
+
+Documentation
+=============
+
+Documentation is online under https://braindecode.org, both in the stable and dev versions.
+
+
+Contributing to Braindecode
+===========================
+
+Guidelines for contributing to the library can be found on the braindecode github:
+
+https://github.com/braindecode/braindecode/blob/master/CONTRIBUTING.md
+
+Braindecode chat
+================
+
+https://gitter.im/braindecodechat/community
+
+Citing
+======
+
+If you use this code in a scientific publication, please cite us as:
+
+.. code-block:: bibtex
+
+  @article {HBM:HBM23730,
+  author = {Schirrmeister, Robin Tibor and Springenberg, Jost Tobias and Fiederer,
+    Lukas Dominique Josef and Glasstetter, Martin and Eggensperger, Katharina and Tangermann, Michael and
+    Hutter, Frank and Burgard, Wolfram and Ball, Tonio},
+  title = {Deep learning with convolutional neural networks for EEG decoding and visualization},
+  journal = {Human Brain Mapping},
+  issn = {1097-0193},
+  url = {http://dx.doi.org/10.1002/hbm.23730},
+  doi = {10.1002/hbm.23730},
+  month = {aug},
+  year = {2017},
+  keywords = {electroencephalography, EEG analysis, machine learning, end-to-end learning, brain–machine interface,
+    brain–computer interface, model interpretability, brain mapping},
+  }
+
+as well as the `MNE-Python <https://mne.tools>`_ software that is used by braindecode:
+
+.. code-block:: bibtex
+
+  @article{10.3389/fnins.2013.00267,
+  author={Gramfort, Alexandre and Luessi, Martin and Larson, Eric and Engemann, Denis and Strohmeier, Daniel and Brodbeck, Christian and Goj, Roman and Jas, Mainak and Brooks, Teon and Parkkonen, Lauri and Hämäläinen, Matti},
+  title={{MEG and EEG data analysis with MNE-Python}},
+  journal={Frontiers in Neuroscience},
+  volume={7},
+  pages={267},
+  year={2013},
+  url={https://www.frontiersin.org/article/10.3389/fnins.2013.00267},
+  doi={10.3389/fnins.2013.00267},
+  issn={1662-453X},
+  }
+
+
+
+
+Licensing
+^^^^^^^^^
+
+Braindecode is **BSD-licenced** (BSD-3-Clause):
+
+    This software is OSI Certified Open Source Software.
+    OSI Certified is a certification mark of the Open Source Initiative.
+
+    Copyright (c) 2011-2022, authors of Braindecode.
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions are met:
+
+    * Redistributions of source code must retain the above copyright notice,
+      this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above copyright notice,
+      this list of conditions and the following disclaimer in the documentation
+      and/or other materials provided with the distribution.
+
+    * Neither the names of braindecode authors nor the names of any
+      contributors may be used to endorse or promote products derived from
+      this software without specific prior written permission.
+
+    **This software is provided by the copyright holders and contributors
+    "as is" and any express or implied warranties, including, but not
+    limited to, the implied warranties of merchantability and fitness for
+    a particular purpose are disclaimed. In no event shall the copyright
+    owner or contributors be liable for any direct, indirect, incidental,
+    special, exemplary, or consequential damages (including, but not
+    limited to, procurement of substitute goods or services; loss of use,
+    data, or profits; or business interruption) however caused and on any
+    theory of liability, whether in contract, strict liability, or tort
+    (including negligence or otherwise) arising in any way out of the use
+    of this software, even if advised of the possibility of such
+    damage.**
+
+
+

{Braindecode-0.6 → Braindecode-0.7}/README.rst

@@ -1,6 +1,3 @@
-Braindecode
-===========
-
 .. image:: https://badges.gitter.im/braindecodechat/community.svg
    :alt: Join the chat at https://gitter.im/braindecodechat/community
    :target: https://gitter.im/braindecodechat/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge
@@ -16,15 +13,23 @@ Braindecode
    :target: https://codecov.io/gh/braindecode/braindecode
    :alt: Code Coverage
 
-A deep learning toolbox to decode raw time-domain EEG.
+.. |Braindecode| image:: https://user-images.githubusercontent.com/42702466/177958779-b00628aa-9155-4c51-96d1-d8c345aff575.svg
+.. _braindecode: braindecode.org/
+
+Braindecode
+===========
+
+Braindecode is an open-source Python toolbox for decoding raw electrophysiological brain
+data with deep learning models. It includes dataset fetchers, data preprocessing and
+visualization tools, as well as implementations of several deep learning
+architectures and data augmentations for analysis of EEG, ECoG and MEG.
 
-For EEG researchers who want to work with deep learning and
-deep learning researchers who want to work with EEG data.
-For now focused on convolutional networks.
+For neuroscientists who want to work with deep learning and
+deep learning researchers who want to work with neurophysiological data.
 
 
-Installation
-============
+Installation Braindecode
+========================
 
 1. Install pytorch from http://pytorch.org/ (you don't need to install torchvision).
 
@@ -58,14 +63,20 @@ alternatively, install the latest version of braindecode via pip:
 Documentation
 =============
 
-Documentation is online under https://braindecode.org
+Documentation is online under https://braindecode.org, both in the stable and dev versions.
+
+
+Contributing to Braindecode
+===========================
 
+Guidelines for contributing to the library can be found on the braindecode github:
 
-Dataset
-=======
-The high-gamma dataset used in our publication (see below), including trained models, is available under:
-https://web.gin.g-node.org/robintibor/high-gamma-dataset/
+https://github.com/braindecode/braindecode/blob/master/CONTRIBUTING.md
 
+Braindecode chat
+================
+
+https://gitter.im/braindecodechat/community
 
 Citing
 ======
@@ -104,3 +115,45 @@ as well as the `MNE-Python <https://mne.tools>`_ software that is used by braind
   doi={10.3389/fnins.2013.00267},
   issn={1662-453X},
   }
+
+
+
+
+Licensing
+^^^^^^^^^
+
+Braindecode is **BSD-licenced** (BSD-3-Clause):
+
+    This software is OSI Certified Open Source Software.
+    OSI Certified is a certification mark of the Open Source Initiative.
+
+    Copyright (c) 2011-2022, authors of Braindecode.
+    All rights reserved.
+
+    Redistribution and use in source and binary forms, with or without
+    modification, are permitted provided that the following conditions are met:
+
+    * Redistributions of source code must retain the above copyright notice,
+      this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above copyright notice,
+      this list of conditions and the following disclaimer in the documentation
+      and/or other materials provided with the distribution.
+
+    * Neither the names of braindecode authors nor the names of any
+      contributors may be used to endorse or promote products derived from
+      this software without specific prior written permission.
+
+    **This software is provided by the copyright holders and contributors
+    "as is" and any express or implied warranties, including, but not
+    limited to, the implied warranties of merchantability and fitness for
+    a particular purpose are disclaimed. In no event shall the copyright
+    owner or contributors be liable for any direct, indirect, incidental,
+    special, exemplary, or consequential damages (including, but not
+    limited to, procurement of substitute goods or services; loss of use,
+    data, or profits; or business interruption) however caused and on any
+    theory of liability, whether in contract, strict liability, or tort
+    (including negligence or otherwise) arising in any way out of the use
+    of this software, even if advised of the possibility of such
+    damage.**
+

{Braindecode-0.6 → Braindecode-0.7}/braindecode/augmentation/base.py

@@ -1,6 +1,7 @@
 # Authors: Cédric Rommel <cedric.rommel@inria.fr>
 #          Alexandre Gramfort <alexandre.gramfort@inria.fr>
-#
+#          Bruno Aristimunha <b.aristimunha@gmail.com>
+#          Martin Wimpff <martin.wimpff@iss.uni-stuttgart.de>
 # License: BSD (3-clause)
 
 from typing import List, Tuple, Any
@@ -50,7 +51,7 @@ class Transform(torch.nn.Module):
         self._probability = probability
         self.rng = check_random_state(random_state)
 
-    def get_params(self, *batch):
+    def get_augmentation_params(self, *batch):
         return dict()
 
     def forward(self, X: Tensor, y: Tensor = None) -> Output:
@@ -87,7 +88,7 @@ class Transform(torch.nn.Module):
             # Uses the mask to define the output
             out_X[mask, ...], tr_y = self.operation(
                 out_X[mask, ...], out_y[mask],
-                **self.get_params(out_X[mask, ...], out_y[mask])
+                **self.get_augmentation_params(out_X[mask, ...], out_y[mask])
             )
             # Apply the operation defining the Transform to the whole batch
             if type(tr_y) is tuple:
@@ -142,11 +143,20 @@ class Compose(Transform):
         return X, y
 
 
-def _make_collateable(transform):
+def _make_collateable(transform, device=None):
+    """ Wraps a transform to make it collateable.
+        with device control. """
+
     def _collate_fn(batch):
         collated_batch = default_collate(batch)
         X, y = collated_batch[:2]
+
+        if device is not None:
+            X = X.to(device)
+            y = y.to(device)
+
         return (*transform(X, y), *collated_batch[2:])
+
     return _collate_fn
 
 
@@ -159,24 +169,26 @@ class AugmentedDataLoader(DataLoader):
         The dataset containing the signals.
     transforms : list | Transform, optional
         Transform or sequence of Transform to be applied to each batch.
+    device : str | torch.device | None, optional
+        Device on which to transform the data. Defaults to None.
     **kwargs : dict, optional
         keyword arguments to pass to standard DataLoader class.
     """
 
-    def __init__(self, dataset, transforms=None, **kwargs):
+    def __init__(self, dataset, transforms=None, device=None, **kwargs):
         if "collate_fn" in kwargs:
             raise ValueError(
                 "collate_fn cannot be used in this context because it is used "
                 "to pass transform"
             )
         if transforms is None or (
-            isinstance(transforms, list) and len(transforms) == 0
+                isinstance(transforms, list) and len(transforms) == 0
         ):
-            self.collated_tr = _make_collateable(IdentityTransform())
+            self.collated_tr = _make_collateable(IdentityTransform(), device=device)
         elif isinstance(transforms, (Transform, nn.Module)):
-            self.collated_tr = _make_collateable(transforms)
+            self.collated_tr = _make_collateable(transforms, device=device)
         elif isinstance(transforms, list):
-            self.collated_tr = _make_collateable(Compose(transforms))
+            self.collated_tr = _make_collateable(Compose(transforms), device=device)
         else:
             raise TypeError("transforms can be either a Transform object" +
                             " or a list of Transform objects.")