[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

examples/example.py

@@ -1,4 +1,10 @@
-from gpu_bartender import VRAMCalculator, ModelArgs, DataArgs, OptimizerArgs, FinetuningArgs
+from gpu_bartender import (
+    DataArgs,
+    FinetuningArgs,
+    ModelArgs,
+    OptimizerArgs,
+    VRAMCalculator,
+)
 
 # Example usage
 model_args = ModelArgs(

gpu_bartender/__init__.py

@@ -1,13 +1,13 @@
 from .calculator import VRAMCalculator
-from .model_args import ModelArgs
+from .data_args import DataArgs
 from .finetuning_args import FinetuningArgs
+from .model_args import ModelArgs
 from .optimizer_args import OptimizerArgs
-from .data_args import DataArgs
 
 __all__ = [
     'VRAMCalculator',
     'ModelArgs',
     'FinetuningArgs',
     'OptimizerArgs',
     'DataArgs'
-]
+]

gpu_bartender/calculator.py

@@ -1,18 +1,19 @@
-from dataclasses import dataclass, field
 from typing import Optional
+
 from .data_args import DataArgs
 from .finetuning_args import FinetuningArgs
 from .model_args import ModelArgs
 from .optimizer_args import OptimizerArgs
 
+
 class VRAMCalculator:
     def __init__(
-        self, 
-        model_args: ModelArgs, 
-        finetuning_args: FinetuningArgs, 
-        optimizer_args: OptimizerArgs, 
-        data_args: DataArgs, 
-        num_gpus: int = 1, 
+        self,
+        model_args: ModelArgs,
+        finetuning_args: FinetuningArgs,
+        optimizer_args: OptimizerArgs,
+        data_args: DataArgs,
+        num_gpus: int = 1,
         unit: str = "MiB"
     ):
         self.model_args = model_args

gpu_bartender/data_args.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 
+
 @dataclass
 class DataArgs:
     batch_size: int = field(default=4)

gpu_bartender/finetuning_args.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
+
 @dataclass
 class LoraArgs:
     lora_alpha: Optional[int] = field(default=None)

gpu_bartender/model_args.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field
 
+
 @dataclass
 class ModelArgs:
     num_params: int = field(default=1)
@@ -8,4 +9,4 @@ class ModelArgs:
     num_attention_heads: int = field(default=1)
     num_key_value_heads: int = field(default=1)
     intermediate_size: int = field(default=1)
-    num_layers: int = field(default=1)
+    num_layers: int = field(default=1)

gpu_bartender/optimizer_args.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Optional
 
+
 @dataclass
 class OptimizerArgs:
     optimizer: str = field(default="adam")