fix(tensor_dict): use get_args for TDCompatible type checking

In Python 3.8+, `isinstance` with a `typing.Union` like `TDCompatible` raises a `TypeError`. The correct way to check against a type alias that is a union is to use `get_args` from the `typing` module.

This change updates the type check to `isinstance(value, get_args(TDCompatible))` to ensure correct runtime behavior across supported Python versions.

Author: Tim Joseph

src/tensorcontainer/tensor_dict.py

@@ -27,6 +27,7 @@
     Union,
     cast,
     overload,
+    get_args
 )
 
 import torch
@@ -177,7 +178,7 @@ def _pytree_flatten(
         metadata: Dict[str, Any] = {}
 
         for key, value in self.data.items():
-            if isinstance(value, TDCompatible):
+            if isinstance(value, get_args(TDCompatible)):
                 td_compatible_leaves.append(value)
                 td_compatible_keys.append(key)
             else: