Fix mypy

pytorch-labs · Feb 4, 2025 · dfc9853 · dfc9853
1 parent d991158
commit dfc9853
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/torchfix/visitors/misc/__init__.py b/torchfix/visitors/misc/__init__.py
@@ -190,11 +190,14 @@ def visit_Call(self, node):
                         dim_arg = self.get_specific_arg(
                             node.args[0].value, arg_name="dim", arg_pos=1
                         )
-
-                        if dim_arg is not None and dim_arg.value.value != "None":
-                            self.add_violation(
-                                node,
-                                error_code=self.ERRORS[0].error_code,
-                                message=self.ERRORS[0].message(),
-                                replacement=None,
-                            )
+                        if dim_arg is not None:
+                            if not (
+                                isinstance(dim_arg.value, cst.Name)
+                                and dim_arg.value.value == "None"
+                            ):
+                                self.add_violation(
+                                    node,
+                                    error_code=self.ERRORS[0].error_code,
+                                    message=self.ERRORS[0].message(),
+                                    replacement=None,
+                                )