add log for device

neuralmagic · Nov 14, 2023 · 21b41d3 · 21b41d3
1 parent 68342bf
commit 21b41d3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/sparseml/transformers/sparsification/obcq/obcq.py b/src/sparseml/transformers/sparsification/obcq/obcq.py
@@ -74,6 +74,7 @@ def one_shot(
 
     # fallback to cpu if cuda not available
     device = _fallback_to_cpu(device)
+    _LOGGER.info(f"Running one_shot on device {device}")
 
     # Load the configuration from the model path
     config = AutoConfig.from_pretrained(model_path)