updates

Author: sayakpaul

src/diffusers/utils/torch_utils.py

@@ -154,10 +154,8 @@ def apply_freeu(
 def get_torch_cuda_device_capability():
     if torch.cuda.is_available():
         device = torch.device("cuda")
-        gpu_name = torch.cuda.get_device_name(device)
         compute_capability = torch.cuda.get_device_capability(device)
         compute_capability = f"{compute_capability[0]}.{compute_capability[1]}"
-        print(f"{gpu_name=}, {compute_capability=}")
         return float(compute_capability)
     else:
         return None

tests/models/test_modeling_common.py

@@ -1385,6 +1385,7 @@ def test_layerwise_casting(storage_dtype, compute_dtype):
     @require_torch_gpu
     def test_layerwise_casting_memory(self):
         MB_TOLERANCE = 0.2
+        LEAST_COMPUTE_CAPABILITY = 8.0
 
         def reset_memory_stats():
             gc.collect()
@@ -1417,7 +1418,7 @@ def get_memory_usage(storage_dtype, compute_dtype):
         self.assertTrue(fp8_e4m3_bf16_memory_footprint < fp8_e4m3_fp32_memory_footprint < fp32_memory_footprint)
         # NOTE: the following assertion will fail on our CI (running Tesla T4) due to bf16 using more memory than fp32.
         # On other devices, such as DGX (Ampere) and Audace (Ada), the test passes.
-        if compute_capability >= 8.9:
+        if compute_capability >= LEAST_COMPUTE_CAPABILITY:
             self.assertTrue(fp8_e4m3_bf16_max_memory < fp8_e4m3_fp32_max_memory)
         # On this dummy test case with a small model, sometimes fp8_e4m3_fp32 max memory usage is higher than fp32 by a few
         # bytes. This only happens for some models, so we allow a small tolerance.