Repository: speechbrain/speechbrain

Issue:
AMP at inference time
### Describe the bug

I think we should have a separate flag to enable AMP at test time, as the change in https://github.com/speechbrain/speechbrain/pull/2406 is not backward compatible. Typically you want to train in fp16 but test in fp32. Lowering precision at inference can hamper accuracy and cause unexpected behaviors.
For example, I'm saving some samples at test time and with the new code with `--precision=fp16` my pipeline stops working:
```
>>> write_audio("sig_pred.wav", sig_pred, 16000)
>>> RuntimeError: Input tensor has to be one of float32, int32, int16 or uint8 type.
```

### Expected behaviour

An option to set the AMP configuration at inference time.

### To Reproduce

_No response_

### Environment Details

_No response_

### Relevant Log Output

_No response_

### Additional Context

_No response_


Discussion:
@Adel-Moumen, what do you think about that?
I think that we can have a new flag that check if we apply it at test time or not. We can improve the data class https://github.com/speechbrain/speechbrain/blob/develop/speechbrain/core.py#L90 to have new args such as `enable_testing` so that by default `enable_testing=False` and we can turn it to `True` by using a new flag `--amp_testing=True` or smth like that. Wdyt ? 

 CC: @asumagic you might want to be part of this discussion. 
Maybe just an `--eval-precision` flag that defaults to `fp32`?

Resolving patch:
diff --git a/speechbrain/core.py b/speechbrain/core.py
--- a/speechbrain/core.py
+++ b/speechbrain/core.py
@@ -70,6 +70,7 @@
     "compile_using_fullgraph": False,
     "compile_using_dynamic_shape_tracing": False,
     "precision": "fp32",
+    "eval_precision": "fp32",
     "auto_mix_prec": False,
     "bfloat16_mix_prec": False,
     "max_grad_norm": 5.0,
@@ -360,6 +361,12 @@ def parse_arguments(arg_list=None):
         help="This flag enables training with automatic mixed-precision."
         "It can be set to `fp32`, `fp16`, or `bf16`.",
     )
+    parser.add_argument(
+        "--eval_precision",
+        type=str,
+        help="This flag enables inference with automatic mixed-precision."
+        "It can be set to `fp32`, `fp16`, or `bf16`.",
+    )
     parser.add_argument(
         "--auto_mix_prec",
         default=None,
@@ -555,6 +562,8 @@ class and override any methods for which the default behavior does not
             The location for performing computations.
         precision (str)
             One of ``fp32``, ``fp16``, ``bf16``.
+        eval_precision (str)
+            One of ``fp32``, ``fp16``, ``bf16``.
         auto_mix_prec (bool)
             If ``True``, automatic mixed-precision (fp16) is used.
             Activate it only with cuda. Note: this is a
@@ -735,12 +744,14 @@ def __init__(  # noqa: C901
             )
             self.precision = "bf16"
 
-        if self.device == "cpu" and self.precision == "fp16":
+        if self.device == "cpu" and (
+            self.precision == "fp16" or self.eval_precision == "fp16"
+        ):
             raise ValueError(
-                "The option `--precision` is enabled with the value "
-                "fp16. This option is not yet supported on CPU. "
-                "Please use `--precision=bf16` instead to get "
-                "mixed precision on CPU."
+                "The option `--precision` or `--eval_precision` is set to fp16. "
+                "This option is not yet supported on CPU. "
+                "Please use `--precision=bf16` or `--eval_precision=bf16` instead "
+                "to enable mixed precision on CPU."
             )
 
         gradscaler_enabled = self.precision == "fp16" and "cuda" in self.device
@@ -1295,7 +1306,7 @@ def evaluate_batch(self, batch, stage):
         -------
         detached loss
         """
-        amp = AMPConfig.from_name(self.precision)
+        amp = AMPConfig.from_name(self.eval_precision)
         if self.use_amp:
             with torch.autocast(
                 dtype=amp.dtype, device_type=torch.device(self.device).type,


Test patch: