improved torch generator exception

jkiesele · Nov 7, 2024 · 65c399b · 65c399b
1 parent 06f6a05
commit 65c399b
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/src/djcdata/torch_interface.py b/src/djcdata/torch_interface.py
@@ -67,9 +67,12 @@ def __iter__(self):
         return self
 
     def __next__(self):
-        data = next(self.iterator)
-        # Convert numpy arrays to torch tensors and move to device
-        return self._convert_to_tensors(data)
+        try:
+            data = next(self.iterator)
+            # Convert numpy arrays to torch tensors and move to device
+            return self._convert_to_tensors(data)
+        except StopIteration:
+            return None
 
     def _convert_to_tensors(self, data):
         # data can be (x, y) or (x, y, w)