FIX Preserve defaultdicts in to_device

Author: aswanth-07

CHANGES.md

@@ -16,6 +16,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix moving `defaultdict` inputs to a device.
+
 ## [1.4.0]
 
 ### Added

skorch/tests/test_utils.py

@@ -1,5 +1,6 @@
 """Test for utils.py"""
 
+from collections import defaultdict
 from copy import deepcopy
 from unittest.mock import patch
 
@@ -374,6 +375,18 @@ def test_check_device_dict_torch_tensor(
         for k in x_dict:
             assert np.allclose(x_dict[k], original_x_dict[k])
 
+    def test_check_device_defaultdict_torch_tensor(self, to_device, x_dict):
+        x_defaultdict = defaultdict(list, x_dict)
+
+        result = to_device(x_defaultdict, device='cpu')
+
+        assert isinstance(result, defaultdict)
+        assert result.default_factory is list
+        assert result['missing'] == []
+        for key, value in x_dict.items():
+            assert torch.equal(result[key], value)
+            assert result[key].device.type == 'cpu'
+
     @pytest.mark.parametrize('device_from, device_to', [
         ('cpu', 'cpu'),
         ('cpu', 'cuda'),

skorch/utils.py

@@ -4,6 +4,7 @@
 
 """
 
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from enum import Enum
@@ -205,7 +206,10 @@ def to_device(X, device):
 
     if isinstance(X, Mapping):
         # dict-like but not a dict
-        return type(X)({key: to_device(val, device) for key, val in X.items()})
+        mapped = {key: to_device(val, device) for key, val in X.items()}
+        if isinstance(X, defaultdict):
+            return type(X)(X.default_factory, mapped)
+        return type(X)(mapped)
 
     # PackedSequence class inherits from a namedtuple
     if isinstance(X, (tuple, list)) and (type(X) != PackedSequence):