Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -215,11 +215,11 @@ class RoBERTaVisualizer(TransformerVisualizer):
         attn_matrix = attn_matrix.to(torch.float16)
         
         attn_layer = attentions[target_layer].squeeze(0).mean(dim=0)  # [seq, seq]
-
-        print('Computing grad norms')
+        print(9,time.time()-start)
+        start = time.time()
+        #print('Computing grad norms')
         grad_norms_list = []
         for k in range(seq_len):
-            start = time.time()
             scalar = attn_layer[:, k].sum()
             grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
             grad_norms = grad.norm(dim=1)
@@ -228,15 +228,16 @@ class RoBERTaVisualizer(TransformerVisualizer):
             grad_norms = torch.round(grad_norms.unsqueeze(1).float() * 100) / 100
             grad_norms = grad_norms.to(torch.float16)
             
-            start = time.time()
         
             grad_norms_list.append(grad_norms)
         
+        print(10,time.time()-start)
         start = time.time()
-        print(9,time.time()-start)
         grad_matrix = torch.cat(grad_norms_list, dim=1)
         grad_matrix = grad_matrix[:seq_len, :seq_len]
         attn_matrix = attn_matrix[:seq_len, :seq_len]
+        
+        print(11,time.time()-start)
         start = time.time()
         
         attn_matrix = torch.round(attn_matrix.float() * 100) / 100
@@ -244,7 +245,7 @@ class RoBERTaVisualizer(TransformerVisualizer):
 
         grad_matrix = torch.round(grad_matrix.float() * 100) / 100
         grad_matrix = grad_matrix.to(torch.float16)
-        start = time.time()
+        print(12,time.time()-start)