Use torch.inference_mode() for lower memory usage during calibration (#20)

Author: mgoin

auto_fp8/quantize.py

@@ -272,13 +272,22 @@ def quantize_activations(
     cleanup_memory()
 
     # Pass through calibration data to measure activation scales
+<<<<<<< HEAD
     with tqdm.tqdm(
         total=calibration_tokens.shape[0], desc="Calibrating activation scales"
     ) as pbar:
         for row_idx in range(calibration_tokens.shape[0]):
             model(calibration_tokens[row_idx].reshape(1, -1))
             cleanup_memory()
             pbar.update(1)
+=======
+    with torch.inference_mode():
+        with tqdm.tqdm(total=calibration_tokens.shape[0], desc="Calibrating activation scales") as pbar:
+            for row_idx in range(calibration_tokens.shape[0]):
+                model(calibration_tokens[row_idx].reshape(1, -1))
+                cleanup_memory()
+                pbar.update(1)
+>>>>>>> b1c6ad6 (Use `torch.inference_mode()` for lower memory usage during calibration (#20))
 
     # Replace dynamic quantizer observer with StaticLinear for export
     for name, quantizer in model.named_modules():