Update build_mlp.py

build_mlp.py

@@ -192,9 +192,9 @@ class PLoRA(nn.Linear):
     def forward(self, x, im_mask=None):
         B, N, C = x.shape
         x = x.reshape(-1, C)
-        im_mask = im_mask.view(-1)
         res = super().forward(x)
         if im_mask is not None:
+            im_mask = im_mask.view(-1)
             if torch.sum(im_mask) > 0:
                 part_x = x[im_mask]
                 res[im_mask] += self.Plora_B(self.Plora_A(