Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -10,6 +10,8 @@ import torch.autograd.functional as Fgrad
 
 CACHE_DIR  = "/data/hf_cache"
 
+
+
 class RoBERTaVisualizer(TransformerVisualizer):
     def __init__(self, task):
         super().__init__()
@@ -73,6 +75,12 @@ class RoBERTaVisualizer(TransformerVisualizer):
 
 
         self.model.to(self.device)
+        # Force materialization of all layers (avoids meta device errors)
+        with torch.no_grad():
+            dummy_ids = torch.tensor([[0, 1]], device=self.device)
+            dummy_mask = torch.tensor([[1, 1]], device=self.device)
+            _ = self.model(input_ids=dummy_ids, attention_mask=dummy_mask)
+
         self.model.eval()
         self.num_attention_layers = self.model.config.num_hidden_layers
 
@@ -199,14 +207,14 @@ class RoBERTaVisualizer(TransformerVisualizer):
             return attentions_condensed
     
         start = time.time()
-        jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds)
+        jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds).to(torch.float16)
         print(jac.shape)
         jac = jac.norm(dim=-1).squeeze(dim=2)
         print(jac.shape)
         seq_len = jac.shape[0]
         print(seq_len)
         grad_matrices_all = [jac[ii,:,:].tolist() for ii in range(seq_len)]
-        
+
         print(31,time.time()-start)
         attn_matrices_all = []
         for target_layer in range(len(attentions)):