trainer classes now takes care of auto-casting numpy to torch tensors, and setting correct device based on model parameter devices

dalle2_pytorch/train.py

@@ -52,12 +52,17 @@ def groupby_prefix_and_trim(prefix, d):
 def cast_torch_tensor(fn):
     @wraps(fn)
     def inner(model, *args, **kwargs):
-        device = next(model.parameters()).device
+        device = kwargs.pop('_device', next(model.parameters()).device)
+        cast_device = kwargs.pop('_cast_device', True)
+
         kwargs_keys = kwargs.keys()
         all_args = (*args, *kwargs.values())
         split_kwargs_index = len(all_args) - len(kwargs_keys)
         all_args = tuple(map(lambda t: torch.from_numpy(t) if exists(t) and isinstance(t, np.ndarray) else t, all_args))
-        all_args = tuple(map(lambda t: t.to(device) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
+
+        if cast_device:
+            all_args = tuple(map(lambda t: t.to(device) if exists(t) and isinstance(t, torch.Tensor) else t, all_args))
+
         args, kwargs_values = all_args[:split_kwargs_index], all_args[split_kwargs_index:]
         kwargs = dict(tuple(zip(kwargs_keys, kwargs_values)))
 

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.2.33',
+  version = '0.2.34',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',