minor

Author: Your Name

modelopt/torch/quantization/plugins/transformers_trainer.py

@@ -194,8 +194,8 @@ def _restore_modelopt_state_with_weights(self):
 
     def _quantize_model(self):
         """Quantize the model. Restore the quantization state if it exists."""
-        num_samples = min(self.quant_args.calib_size, len(self.eval_dataset))  # type: ignore [union-attr]
         dataset = self.train_dataset if self.train_dataset is not None else self.eval_dataset
+        num_samples = min(self.quant_args.calib_size, len(dataset))  # type: ignore [union-attr]
         dataset = torch.utils.data.Subset(dataset, list(range(num_samples)))
         data_loader = self.get_eval_dataloader(dataset)