Test AugmentedDataLoader device kwarg
@@ -11,6 +11,7 @@ from sklearn.utils import check_random_state
from braindecode.augmentation.base import AugmentedDataLoader
from braindecode.augmentation.base import Compose
from braindecode.augmentation.base import Transform
+from braindecode.augmentation.transforms import SmoothTimeMask
from braindecode.datautil import create_from_mne_epochs
@@ -141,18 +142,22 @@ def concat_windows_dataset():
# test AugmentedDataLoader with 0, 1 and 2 composed transforms
-@pytest.mark.parametrize("nb_transforms,no_list", [
- (0, False), (1, False), (1, True), (2, False)
+@pytest.mark.parametrize("nb_transforms,no_list,dummy", [
+ (0, False, True), (1, False, True), (1, True, True), (2, False, True),
+ (1, False, False)
])
def test_data_loader(dummy_transform, concat_windows_dataset, nb_transforms,
- no_list):
- transforms = [dummy_transform for _ in range(nb_transforms)]
+ no_list, dummy):
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
+ transform = dummy_transform if dummy else SmoothTimeMask(0.5)
+ transforms = [transform for _ in range(nb_transforms)]
if no_list:
transforms = transforms[0]
data_loader = AugmentedDataLoader(
concat_windows_dataset,
transforms=transforms,
- batch_size=128)
+ batch_size=128,
+ device=device)
for idx_batch, _ in enumerate(data_loader):
if idx_batch >= 3:
break
