fixes 550 (#853)

keras-team · Sep 16, 2023 · 6ad426a · 6ad426a
1 parent c663efd
commit 6ad426a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras_core/trainers/data_adapters/torch_data_adapter.py b/keras_core/trainers/data_adapters/torch_data_adapter.py
@@ -22,7 +22,7 @@ def __init__(self, dataloader):
 
     def get_numpy_iterator(self):
         for batch in self._dataloader:
-            yield tuple(tree.map_structure(lambda x: x.numpy(), batch))
+            yield tuple(tree.map_structure(lambda x: x.cpu().numpy(), batch))
 
     def get_torch_dataloader(self):
         return self._dataloader