fix: dataloader

Grapycal · Jun 14, 2024 · 46044ba · 46044ba
1 parent d106ef7
commit 46044ba
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/extensions/grapycal_torch/grapycal_torch/dataloader.py b/extensions/grapycal_torch/grapycal_torch/dataloader.py
@@ -4,6 +4,7 @@
 from grapycal.stores import main_store
 from grapycal_torch.store import GrapycalTorchStore
 from topicsync.topic import GenericTopic
+from torch import Tensor
 from torch.utils.data import DataLoader
 
 
@@ -56,7 +57,18 @@ def task(self):
                     )
                     return
                 if self.to_defalut_device.get():
-                    batch = self.get_store(GrapycalTorchStore).to_default_device(batch)
+                    if isinstance(batch, dict):
+                        for key in batch:
+                            if isinstance(batch[key], Tensor):
+                                batch[key] = self.get_store(
+                                    GrapycalTorchStore
+                                ).to_default_device(batch[key])
+                    elif isinstance(batch, Tensor):
+                        batch = self.get_store(GrapycalTorchStore).to_default_device(
+                            batch
+                        )
+                    else:
+                        pass
                 self.out.push(batch)
                 self.flash_running_indicator()
                 yield

diff --git a/extensions/grapycal_torch/grapycal_torch/moduleNode.py b/extensions/grapycal_torch/grapycal_torch/moduleNode.py
@@ -157,7 +157,7 @@ def destroy(self):
 
 class SimpleModuleNode(ModuleNode):
     module_type: type[nn.Module] = nn.Module
-    inputs: list[str] = ["input"]
+    inputs: list[str] = []
     max_in_degree = [1]
     outputs = ["output"]
     display_port_names: bool | None = None