free extra memory

vllm-project · Aug 13, 2024 · 453022a · 453022a
1 parent b7280e1
commit 453022a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/llmcompressor/modifiers/smoothquant/base.py b/src/llmcompressor/modifiers/smoothquant/base.py
@@ -305,6 +305,9 @@ def smooth(module):
                     smooth(layer)
                 smooth(smooth_layer)
 
+        # clear out allocated smoothing scales
+        torch.cuda.empty_cache()
+
     def _calculate_smoothing_scales(
         self, balance_layers: List[Module], activation_scales: torch.Tensor
     ) -> List[float]: