Spaces:

RWKV-Red-Team
/

RWKV-LatestSpace

Running on T4

sparkleman commited on Mar 14

Commit

d055fef

1 Parent(s): 9b15eee

UPDATE: Torch memory

Files changed (1) hide show

app.py CHANGED Viewed

@@ -39,7 +39,7 @@ def logGPUState():
     if "cuda" in CONFIG.STRATEGY:
         gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
         logger.info(
-            f"[STATUS] vram {format_bytes(gpu_info.total)} used {format_bytes(gpu_info.used)} free {format_bytes(gpu_info.free)}"
         )

     if "cuda" in CONFIG.STRATEGY:
         gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
         logger.info(
+            f"[STATUS] Torch - {format_bytes(torch.cuda.memory_allocated())} - NVML - vram {format_bytes(gpu_info.total)} used {format_bytes(gpu_info.used)} free {format_bytes(gpu_info.free)}"
         )