fix tests

Author: drisspg

float8_experimental/float8_tensor.py

@@ -266,9 +266,13 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None):
 
         # All ops in the FLOAT8_OPS_TABLE expect Float8Tensor as inputs
         # And don't support mixed tensor subclasses. This will trigger the handler for
-        # the next type in the dispatch list. torch._C._TensorMeta is for FakeTensor
+        # the next type in the dispatch list
         def allowed_subclasses(type):
-            return issubclass(cls, type) or isinstance(type, torch._C._TensorMeta)
+            return (
+                issubclass(cls, type)
+                or issubclass(torch._subclasses.fake_tensor.FakeTensor, type)
+                or issubclass(torch._subclasses.functional_tensor.FunctionalTensor, type)
+            )
 
         if not all(allowed_subclasses(t) for t in types):
             return NotImplemented

test/test_everything.sh

@@ -9,5 +9,6 @@ pytest test/test_compile.py
 ./test/test_fsdp.sh
 ./test/test_fsdp_compile.sh
 ./test/test_tp.sh
+./test/test_dtensor.sh
 
 echo "all tests successful"