align dim size funtion (#4880)

* align dim size funtion

* fix dim usage
Co-authored-by: Noneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

oneflow/python/framework/tensor.py

@@ -222,11 +222,14 @@ class Tensor:
         else:
             self._undetermined_tensor.requires_grad = requires_grad
 
-    def size(self):
-        return self.shape
+    def size(self, idx=None):
+        if idx is None:
+            return self.shape
+        else:
+            return self.shape[idx]
 
-    def dim(self, idx):
-        return self.shape[idx]
+    def dim(self):
+        return self.ndim
 
     def ndimension(self):
         return self.ndim

oneflow/python/nn/init.py

@@ -79,8 +79,8 @@ def _calculate_fan_in_and_fan_out(tensor):
             "Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
         )
 
-    num_input_fmaps = tensor.dim(1)
-    num_output_fmaps = tensor.dim(0)
+    num_input_fmaps = tensor.size(1)
+    num_output_fmaps = tensor.size(0)
     receptive_field_size = 1
     if tensor.ndimension() > 2:
         # math.prod is not always available, accumulate the product manually