Update classifier.py

classifier.py

@@ -95,8 +95,8 @@ class Classifier:
             grad = []
             # Normalize the weight before inference
             # This will constrain the gradient or you will have an explosion on query vector
-            self.weight.data /= torch.norm(
-                self.weight.data, p=2, dim=-1, keepdim=True
+            self.weight /= torch.norm(
+                self.weight, p=2, dim=-1, keepdim=True
             )
             for n in range(self.num_class):
                 # select all training sample and create labels
@@ -119,6 +119,9 @@ class Classifier:
             # update weights
             grad = torch.stack(grad, dim=0)
             self.weight -= 0.1 * grad
+            self.weight /= torch.norm(
+                self.weight, p=2, dim=-1, keepdim=True
+            )
 
     def get_weights(self):
         xq = self.weight.detach().numpy()