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,17 +52,12 @@ def groupby_prefix_and_trim(prefix, d):
 def cast_torch_tensor(fn):
     @wraps(fn)
     def inner(model, *args, **kwargs):
-        device = kwargs.pop('_device', next(model.parameters()).device)
+        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.34',
+  version = '0.2.33',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',