Fix: process torch.device and torch.dtype in inputs

Author: JewelRoam

graph_net/torch/decompose_util.py

@@ -248,8 +248,14 @@ def get_args_node(arg):
             yield arg.start
             yield arg.stop
             yield arg.step
+        elif isinstance(arg, torch.device):
+            pass
+        elif isinstance(arg, torch.dtype):
+            pass
         else:
-            assert isinstance(arg, (int, bool, float, str, type(None))), f"{type(arg)=}"
+            assert isinstance(
+                arg, (int, bool, float, str, type(None), torch.device, torch.dtype)
+            ), f"{type(arg)=}"
 
     def get_args_node_and_self_node(node):
         for arg in node.args: