make cleanup function public

README.md

@@ -101,6 +101,14 @@ attributor = Attributor(
 )
 ```
 
+### Cleaning Up
+
+A convenience method is provided to clean up memory used by Python and Torch. This can be useful when running the library in a cloud notebook environment:
+
+```python
+attributor.cleanup()
+```
+
 ## Development
 
 To contribute to the library, you will need to install the development requirements:

attribution/attribution.py

@@ -80,8 +80,6 @@ def get_attributions(
             attr_scores[it] = attr_scores_next_token
             token_ids = torch.cat((token_ids, next_token_id.view(-1)), dim=0)
 
-            self._cleanup()
-
         return attr_scores, token_ids
 
     def print_attributions(
@@ -106,6 +104,20 @@ def print_attributions(
             word_list, attr_scores, token_ids, generation_length
         )
 
+    def cleanup(self) -> None:
+        """
+        This function is used to free up the memory resources. It clears the GPU cache and triggers garbage collection.
+
+        Returns:
+        None
+        """
+        if hasattr(torch, self.device) and hasattr(
+            getattr(torch, self.device), "empty_cache"
+        ):
+            logging.info(f"Clearing {self.device} cache")
+            getattr(torch, self.device).empty_cache()
+        gc.collect()
+
     def _get_input_embeddings(
         self, embeddings: torch.Tensor, token_ids: torch.Tensor
     ) -> torch.Tensor:
@@ -152,10 +164,6 @@ def _get_attr_scores_next_token(
             attr_scores_next_token[i] = presence_grad
         return attr_scores_next_token
 
-    def _cleanup(self) -> None:
-        torch.cuda.empty_cache()
-        gc.collect()
-
     def _validate_inputs(
         self,
         model: transformers.GPT2LMHeadModel,