Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -193,17 +193,20 @@ class RoBERTaVisualizer(TransformerVisualizer):
                 attention_mask=inputs["attention_mask"],
                 output_attentions=True
             )
+            attentions = outputs.attentions
             attentions_condensed = [a.mean(dim=0).mean(dim=0).sum(dim=0) for a in attentions]
-            print([a.shape for a in attentions_condensed])
             attentions_condensed= torch.vstack(attentions_condensed)
-            print(attentions_condensed.shape) 
             return attentions_condensed
     
         start = time.time()
-        jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds).norm(dim=-1).squeeze(dim=2)
+        jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds)
         print(jac.shape)
-        grad_matrices_all = [jac[i] for i in range(jac.size(0))]
-
+        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)):
@@ -212,7 +215,9 @@ class RoBERTaVisualizer(TransformerVisualizer):
             attn_matrix = mean_attns[target_layer]
             seq_len = attn_matrix.shape[0]
             attn_matrix = attn_matrix[:seq_len, :seq_len]
+            print(4,attn_matrix.shape)
             attn_matrices_all.append(attn_matrix.tolist())
+
         print(3,time.time()-start)
  
 
@@ -259,6 +264,6 @@ class RoBERTaVisualizer(TransformerVisualizer):
                 grad_norms_list.append(grad_norms.unsqueeze(1))
             print(2,time.time()-start)
         """
-
+        #print(grad_matrices_all)
         return grad_matrices_all, attn_matrices_all