linter stuff

maximegmd · Feb 4, 2024 · 67c70d1 · 67c70d1
1 parent eb300b6
commit 67c70d1
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 2 deletions.
diff --git a/setup.py b/setup.py
@@ -1,8 +1,9 @@
 """setup.py for axolotl"""
 
+import platform
 from importlib.metadata import PackageNotFoundError, version
+
 from packaging.version import Version, parse
-import platform
 from setuptools import find_packages, setup
 
 

diff --git a/src/axolotl/utils/bench.py b/src/axolotl/utils/bench.py
@@ -46,6 +46,7 @@ def gpu_memory_usage_all(device=0):
     smi = gpu_memory_usage_smi(device)
     return usage, reserved - usage, max(0, smi - reserved)
 
+
 def mps_memory_usage_all():
     usage = torch.mps.current_allocated_memory() / 1024.0**3
     reserved = torch.mps.driver_allocated_memory() / 1024.0**3

diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -672,7 +672,7 @@ def load_model(
     ):
         model.config.eos_token_id = tokenizer.eos_token_id
 
-    if hasattr(model, "device") and (model.device.type == "cuda" or model.device.type == "mps"):
+    if hasattr(model, "device") and model.device.type in ("cuda", "mps"):
         log_gpu_memory_usage(LOG, "after model load", model.device)
 
     # make sure these are fp32 per Ramesh et al. (2021)