expose VRAM easily

tortoise/utils/device.py

@@ -47,19 +47,19 @@ def get_device(verbose=False):
 
     return torch.device(name)
 
-def get_device_batch_size():
+def get_device_vram( name=get_device_name() ):
     available = 1
-    name = get_device_name()
 
-    if name == "dml":
-        # there's nothing publically accessible in the DML API that exposes this
-        # there's a method to get currently used RAM statistics... as tiles
-        available = 1
-    elif name == "cuda":
+    if name == "cuda":
         _, available = torch.cuda.mem_get_info()
     elif name == "cpu":
         available = psutil.virtual_memory()[4]
 
+    return available
+
+def get_device_batch_size(name=None):
+    available = get_device_vram(name)
+
     vram = available / (1024 ** 3)
     # I'll need to rework this better
     # simply adding more tiers clearly is not a good way to go about it