correct implementation of ndim in order to work in compile mode with …

…tensorflow (see issue tensorflow/tensorflow#48612 and jonasrauber#36)
eserie · Apr 21, 2021 · 7ce3ceb · 7ce3ceb
1 parent 4656839
commit 7ce3ceb
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/eagerpy/tensor/base.py b/eagerpy/tensor/base.py
@@ -112,4 +112,4 @@ def __pow__(self: TensorType, exponent: TensorOrScalar) -> TensorType:
     @final
     @property
     def ndim(self: TensorType) -> int:
-        return cast(int, self.raw.ndim)
+        return len(self.raw.shape)