[Do Not Merge] Add Pytorch memory logging to crossfit

crossfit/backend/torch/op/base.py

@@ -54,6 +54,12 @@ def __init__(
 
     @torch.no_grad()
     def call(self, data, partition_info=None):
+        # Get the current CUDA device
+        current_device = torch.cuda.current_device()
+        # Print CUDA memory at the beginning of the method
+        print(f"CUDA memory at start (device {current_device}):")
+        print(torch.cuda.memory_summary(device=current_device))
+
         index = data.index.copy()
         if self.sorted_data_loader:
             loader = SortedSeqLoader(
@@ -101,6 +107,10 @@ def call(self, data, partition_info=None):
             raise RuntimeError(f"Unexpected output shape: {output.shape}")
         del outputs, _index
         cleanup_torch_cache()
+
+        # Print CUDA memory at the end of the method
+        print(f"CUDA memory at end (device {current_device}):")
+        print(torch.cuda.memory_summary(device=current_device))
         return out
 
     def meta(self):

crossfit/data/dataframe/core.py

@@ -348,7 +348,7 @@ def _(data):
 
 # Fall-back `ArrayBundle` definition
 class ArrayBundle(FrameBackend):
-    @lru_cache
+    @lru_cache(maxsize=1)
     def __len__(self):
         _len = None
         for k, v in self.data.items():
@@ -421,7 +421,7 @@ def assign(self, **kwargs):
         data = self.data.copy()
         for k, v in kwargs.items():
             if self.columns and len(v) != len(self):
-                raise ValueError(f"Column {k} was length {len(v)}, but expected length {len(self)}")
+                raise ValueError(f"Column {k} was length {len(v)} but expected length {len(self)}")
         data.update(**kwargs)
         return self.__class__(data)