Remove: export to onnx by default (#591)

* Remove: `export` to `onnx` by default

* Remove: not needed params

Author: rahul-tuli

src/sparseml/pytorch/image_classification/train.py

@@ -615,7 +615,6 @@ def train(
                     helpers.save_model_training(
                         model,
                         optim,
-                        input_shape,
                         "checkpoint-best",
                         save_dir,
                         epoch,
@@ -633,7 +632,6 @@ def train(
                 helpers.save_model_training(
                     model,
                     optim,
-                    input_shape,
                     f"checkpoint-{epoch:04d}-{val_metric:.04f}",
                     save_dir,
                     epoch,
@@ -648,7 +646,7 @@ def train(
             # only convert qat -> quantized ONNX graph for finalized model
             # TODO: change this to all checkpoints when conversion times improve
             helpers.save_model_training(
-                model, optim, input_shape, "model", save_dir, epoch - 1, val_res, True
+                model, optim, "model", save_dir, epoch - 1, val_res
             )
 
             LOGGER.info("layer sparsities:")

src/sparseml/pytorch/image_classification/utils/helpers.py

@@ -378,12 +378,10 @@ def save_recipe(
 def save_model_training(
     model: Module,
     optim: Optimizer,
-    input_shape: Tuple[int, ...],
     save_name: str,
     save_dir: str,
     epoch: int,
     val_res: Union[ModuleRunResults, None],
-    convert_qat: bool = False,
 ):
     """
     :param model: model architecture
@@ -404,12 +402,6 @@ def save_model_training(
     )
     exporter = ModuleExporter(model, save_dir)
     exporter.export_pytorch(optim, epoch, f"{save_name}.pth")
-    exporter.export_onnx(
-        torch.randn(1, *input_shape),
-        f"{save_name}.onnx",
-        convert_qat=convert_qat,
-    )
-
     info_path = os.path.join(save_dir, f"{save_name}.txt")
 
     with open(info_path, "w") as info_file: