Update modeling_llama.py

modeling_llama.py

@@ -457,7 +457,7 @@ class LlamaAttention(nn.Module):
         if attn_weights.shape[2]>576:
             # print("loading ... ")
             #print(value_states.shape)
-            self.ae_v.load_state_dict(torch.load("weights_320/"+"autoencoder_epoch_1_L1_1280_nonorm_layer_"+str(self.layer_idx)+".pth", map_location='cuda'))
+            self.ae_v.load_state_dict(torch.load("weights_320/"+"autoencoder_epoch_1_L1_nonorm_layer_"+str(self.layer_idx)+".pth", map_location='cuda'))
             value_states_v = value_states[:,:,35:35+576,:]
             value_states_v = value_states_v.permute(0, 2, 1, 3)
             value_states_v=value_states_v.reshape(value_states_v.shape[0],value_states_v.shape[1],5120)