[Refactor] Replace and ban a few torch.cuda functions in favor of torch.accelerator replacements. (#3365)

Signed-off-by: Nick Cao <ncao@redhat.com>

tests/diffusion/quantization/test_quantization_quality.py

@@ -104,7 +104,7 @@ def _generate_image(omni, config: QualityTestConfig):
     generator = torch.Generator(
         device=current_omni_platform.device_type,
     ).manual_seed(config.seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
         {"prompt": config.prompt},
@@ -116,7 +116,7 @@ def _generate_image(omni, config: QualityTestConfig):
         ),
     )
 
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
     first = outputs[0]
     if hasattr(first, "images") and first.images:
         return first.images[0], peak_mem
@@ -135,7 +135,7 @@ def _generate_video(omni, config: QualityTestConfig):
     generator = torch.Generator(
         device=current_omni_platform.device_type,
     ).manual_seed(config.seed)
-    torch.cuda.reset_peak_memory_stats()
+    torch.accelerator.reset_peak_memory_stats()
 
     outputs = omni.generate(
         {"prompt": config.prompt, "negative_prompt": ""},
@@ -148,7 +148,7 @@ def _generate_video(omni, config: QualityTestConfig):
         ),
     )
 
-    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+    peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
     first = outputs[0]
     if hasattr(first, "request_output") and isinstance(first.request_output, list):
         inner = first.request_output[0]
@@ -190,8 +190,8 @@ def _unload(omni):
     del omni
     gc.collect()
     if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
+        torch.accelerator.empty_cache()
+        torch.accelerator.synchronize()
 
 
 # ---------------------------------------------------------------------------