quality

huggingface · Mar 15, 2024 · 4a92318 · 4a92318
1 parent b83e125
commit 4a92318
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 8 deletions.
diff --git a/src/optimum/nvidia/quantization/ammo/config.py b/src/optimum/nvidia/quantization/ammo/config.py
@@ -140,7 +140,7 @@ def from_description(
         num_samples: int = 512,
         max_sequence_length: int = 1024,
         seed: int = 2016,
-        device: Union[str, torch.device] = "cpu"
+        device: Union[str, torch.device] = "cpu",
     ):
         random.seed(seed)
         np.random.seed(seed)
@@ -160,7 +160,7 @@ def from_description(
                 seqlen=max_sequence_length,
                 split=split,
                 seed=seed,
-                device=device
+                device=device,
             )
         else:
             raise ValueError("Providing custom dataset is not yet supported")

diff --git a/src/optimum/nvidia/quantization/ammo/quantizer.py b/src/optimum/nvidia/quantization/ammo/quantizer.py
@@ -118,8 +118,7 @@ def _loop():
 
                     for sample in tqdm(data):
                         inputs = {
-                            name: tensor.to("cuda:0")
-                            for name, tensor in sample.items()
+                            name: tensor.to("cuda:0") for name, tensor in sample.items()
                         }
                         model(**inputs)
 

diff --git a/src/optimum/nvidia/quantization/datasets.py b/src/optimum/nvidia/quantization/datasets.py
@@ -216,12 +216,13 @@ def get_c4_new(
         data[indexes]["text"],
         truncation=True,
         max_length=seqlen,
-        return_attention_mask=False
+        return_attention_mask=False,
     )
 
-    dataset = [{
-        "input_ids": torch.tensor(tokens.ids, dtype=torch.long, device=device)
-    } for tokens in encodings.encodings]
+    dataset = [
+        {"input_ids": torch.tensor(tokens.ids, dtype=torch.long, device=device)}
+        for tokens in encodings.encodings
+    ]
 
     return dataset