[fx] added ndim property to proxy (#1253)

2022-07-12 15:27:13 +08:00 · 2022-07-12 15:27:13 +08:00 · fb35460595
parent 4a09fc0947
commit fb35460595
1 changed files with 4 additions and 0 deletions
--- a/colossalai/fx/proxy.py
+++ b/colossalai/fx/proxy.py
@ -57,6 +57,10 @@ class ColoProxy(Proxy):
        self._assert_meta_data_is_tensor()
        return self.meta_data.shape

+    @property
+    def ndim(self):
+        return self.dim()
+
    def dim(self):
        self._assert_meta_data_is_tensor()
        return self.meta_data.dim()