refine CUDA Graph (#38401)

python/paddle/device/cuda/graphs.py

@@ -17,9 +17,14 @@ from paddle.fluid.core import is_compiled_with_cuda, is_compiled_with_rocm, CUDA
 
 if is_compiled_with_cuda() and not is_compiled_with_rocm():
     from paddle.fluid.core import CUDAGraph as CoreCUDAGraph
+else:
+    CoreCUDAGraph = None
+
 
-    class CUDAGraph:
+class CUDAGraph:
     def __init__(self, place=None, mode="thread_local"):
+        assert CoreCUDAGraph is not None, "CUDA Graph is only supported on PaddlePaddle compiled with NVIDIA GPU."
+
         ALL_MODES = ["global", "thread_local", "relaxed"]
         self._graph = None
         if place is None:
@@ -50,23 +55,3 @@ if is_compiled_with_cuda() and not is_compiled_with_rocm():
         if flags is None:
             flags = 2047  # only all information. It can be any integer inside [1, 2048)  
         self._graph.print_to_dot_files(dirname, flags)
-else:
-
-    class CUDAGraph:
-        def __init__(self, place=None, mode="thread_local"):
-            raise NotImplementedError()
-
-        def capture_begin(self):
-            raise NotImplementedError()
-
-        def capture_end(self):
-            raise NotImplementedError()
-
-        def replay(self):
-            raise NotImplementedError()
-
-        def reset(self):
-            raise NotImplementedError()
-
-        def print_to_dot_files(self, dirname, flags=None):
-            raise NotImplementedError()