tmp set the loading device to cuda:0

Author: frankaging

pyvene/models/intervenable_base.py

@@ -1373,7 +1373,8 @@ def load_intervention(self, load_directory, include_model=True):
             intervention = v
             binary_filename = f"intkey_{k}.bin"
             if isinstance(intervention, TrainableIntervention):
-                saved_state_dict = torch.load(os.path.join(load_directory, binary_filename))
+                saved_state_dict = torch.load(
+                    os.path.join(load_directory, binary_filename), map_location='cuda:0')
                 intervention.load_state_dict(saved_state_dict)
 
         # load model's trainable parameters as well