Fix pin_memory bug

yzhangcs · yzhangcs · commit 2c0eff6e64d3 · 2022-06-16T18:31:11.000+08:00
diff --git a/supar/utils/transform.py b/supar/utils/transform.py
@@ -697,8 +697,8 @@ def compose(self, transform: Transform):
         return [f.compose([s.fields[f.name] for s in self.sentences]) for f in transform.flattened_fields]
 
     def pin_memory(self):
-        for name in self.names:
-            for i in getattr(self, name):
+        for s in self.sentences:
+            for i in s.fields.values():
                 if isinstance(i, torch.Tensor):
                     i.pin_memory()
         return self
