Update lora_w2w.py

lora_w2w.py

@@ -71,13 +71,13 @@ class LoRAModule(nn.Module):
         self.in_dim = org_module.in_features
         self.out_dim = org_module.out_features
 
-        self.proj = proj.bfloat16()
-        self.mean1 = mean[0:self.in_dim].bfloat16()
-        self.mean2 = mean[self.in_dim:].bfloat16()
-        self.std1 = std[0:self.in_dim].bfloat16()
-        self.std2 = std[self.in_dim:].bfloat16()
-        self.v1 = v[0:self.in_dim].bfloat16()
-        self.v2 = v[self.in_dim: ].bfloat16()
+        self.proj = proj.bfloat16().cuda()
+        self.mean1 = mean[0:self.in_dim].bfloat16().cuda()
+        self.mean2 = mean[self.in_dim:].bfloat16().cuda()
+        self.std1 = std[0:self.in_dim].bfloat16().cuda()
+        self.std2 = std[self.in_dim:].bfloat16().cuda()
+        self.v1 = v[0:self.in_dim].bfloat16().cuda()
+        self.v2 = v[self.in_dim: ].bfloat16().cuda()
 
         if type(alpha) == torch.Tensor:
             alpha = alpha.detach().numpy()
@@ -95,6 +95,10 @@ class LoRAModule(nn.Module):
         del self.org_module
 
     def forward(self, x):
+        print(self.proj.device)
+        print(self.v1.device)
+        print(self.mean1.device)
+        print(self.std1.device)
         return self.org_forward(x) +\
             (x@(([email protected])*self.std1+self.mean1).T)@((([email protected])*self.std2+self.mean2))*self.multiplier*self.scale