Bypass non-implemented backward test

pytorch-labs · Nov 18, 2024 · 53a1115 · 53a1115
1 parent 19d6db9
commit 53a1115
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/test/test_gpu/main.py b/test/test_gpu/main.py
@@ -96,10 +96,15 @@ def _run_operator_in_task(op: str, args: List[str]):
     task.check_output()
     task.del_op_instance()
     # Test backward (if applicable)
-    args.extend(["--bwd"])
-    task.make_operator_instance(args=args)
-    task.run()
-    task.check_output()
+    try:
+        args.extend(["--bwd"])
+        task.make_operator_instance(args=args)
+        task.run()
+        task.check_output()
+    except NotImplementedError:
+        logger.info(
+            f"Operator {op.name} does not support backward, skipping backward test."
+        )
 
 
 def make_test(operator):