Merge pull request #150 from ViCCo-Group/mem_leak_fix

Memory leak fix when only extracting cls token
ViCCo-Group · Aug 4, 2023 · 8a3a416 · 8a3a416
2 parents 39dd9d7 + 68ab922
commit 8a3a416
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
@@ -109,7 +109,7 @@ def _extract_batch(
         act = self.activations[module_name]
         if hasattr(self, "extract_cls_token"):
             # we are only interested in the representations of the first token, i.e., [cls] token
-            act = act[:, 0, :]
+            act = act[:, 0, :].clone()
         if flatten_acts:
             if self.model_name.lower().startswith("clip"):
                 act = self.flatten_acts(act, batch, module_name)