revised loading to recycle state dict (#706)

* revised loading to recycle state dict

* removed manuall gc collection

transformer_lens/HookedTransformer.py

@@ -8,7 +8,6 @@
 alteration of activations in individual components like attention heads and MLP layers, facilitating
 a deeper understanding of the internal workings of transformers like GPT-2.
 """
-
 import logging
 import os
 from typing import Dict, List, NamedTuple, Optional, Tuple, Union, cast, overload
@@ -1570,7 +1569,10 @@ def load_and_process_state_dict(
             # so that quantization settings are not lost
             self.load_state_dict(state_dict, assign=True, strict=False)
         else:
-            self.load_state_dict(state_dict, strict=False)
+            state_dict_keys = list(state_dict.keys())
+            for key in state_dict_keys:
+                self.load_state_dict({key: state_dict[key]}, strict=False)
+                del state_dict[key]
 
     def fill_missing_keys(self, state_dict):
         return loading.fill_missing_keys(self, state_dict)