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

fb354605 · Frank Lee · GitHub · 4a09fc09 · fb354605
Unverified Commit fb354605 authored Jul 12, 2022 by Frank Lee Committed by GitHub Jul 12, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 0 deletions

colossalai/fx/proxy.py colossalai/fx/proxy.py +4 -0

No files found.
--- 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()