Update csatv2.py

Author: gusdlf93

timm/models/csatv2.py

@@ -135,7 +135,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x_Y = self.Y_Conv(x[:, 0, ])
         x_Cb = self.Cb_Conv(x[:, 1, ])
         x_Cr = self.Cr_Conv(x[:, 2, ])
-        x = torch.cat([x_Y, x_Cb, x_Cr], axis=1)
+        x = torch.cat([x_Y, x_Cb, x_Cr], dim=1)
         return x
 
 
@@ -160,12 +160,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         mean_list = torch.zeros([3, 64])
         var_list = torch.zeros([3, 64])
-        mean_list[0] = torch.mean(x[:, 0, ], axis=0)
-        mean_list[1] = torch.mean(x[:, 1, ], axis=0)
-        mean_list[2] = torch.mean(x[:, 2, ], axis=0)
-        var_list[0] = torch.var(x[:, 0, ], axis=0)
-        var_list[1] = torch.var(x[:, 1, ], axis=0)
-        var_list[2] = torch.var(x[:, 2, ], axis=0)
+        mean_list[0] = torch.mean(x[:, 0, ], dim=0)
+        mean_list[1] = torch.mean(x[:, 1, ], dim=0)
+        mean_list[2] = torch.mean(x[:, 2, ], dim=0)
+        var_list[0] = torch.var(x[:, 0, ], dim=0)
+        var_list[1] = torch.var(x[:, 1, ], dim=0)
+        var_list[2] = torch.var(x[:, 2, ], dim=0)
         return mean_list, var_list
 
 
@@ -470,6 +470,11 @@ def forward(self, x):
         x = self.forward_head(x)
         return x
 
+
+# --- Components like LayerNorm, GRN, DropPath, FeedForward, PosCNN, trunc_normal_ ---
+# (이 부분은 einops와 무관하므로 위 코드와 동일하게 유지합니다. 여기서는 공간 절약을 위해 생략)
+# 기존 코드의 LayerNorm, GRN, DropPath, FeedForward, PosCNN, trunc_normal_ 함수를 그대로 사용하세요.
+
 class LayerNorm(nn.Module):
     def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
         super().__init__()
@@ -604,5 +609,4 @@ def csatv2(pretrained: bool = False, **kwargs) -> CSATv2:
         img_size=kwargs.pop('img_size', 512),
         num_classes=kwargs.pop('num_classes', 1000),
     )
-
     return _create_csatv2('csatv2', pretrained, **dict(model_args, **kwargs))