Update modelling_uniformer.py

modelling_uniformer.py

@@ -99,8 +99,8 @@ class CBlock(nn.Module):
 
         self.ls = layer_scale
         if self.ls:
-            self.gamma_1 = nn.Parameter(init_value * torch.ones((dim)),requires_grad=True)
-            self.gamma_2 = nn.Parameter(init_value * torch.ones((dim)),requires_grad=True)
+            self.gamma_1 = nn.Parameter(init_value * torch.ones((1, dim, 1, 1)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_value * torch.ones((1, dim, 1, 1)),requires_grad=True)
     
     def forward(self, x):
         x = x + self.pos_embed(x)