Update README.md

README.md

@@ -19,11 +19,6 @@ For example, the file in `gpt2-medium/wikipedia/6_9.pickle` contains the matrix
 import pickle
 import torch
 
-
-def mul(mat, v):
-    return (mat @ v[..., None]).squeeze(-1)
-
-
 with open(file_name, 'rb') as f:
     mat = pickle.load(f)
 
@@ -33,7 +28,7 @@ assert(mat.shape[0] == mat.shape[1])
 
 v = torch.rand(mat.shape[1])
 
-w = mul(mat, v)
+w = mat @ v
 
 assert(w.shape == v.shape)
 ```