evaluate_batch with support of AMP (speechbrain#2406)

* Update core.py * Update core.py
poonehmousavi · Feb 12, 2024 · b8a3ee3 · b8a3ee3
1 parent 64ea6a2
commit b8a3ee3
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/speechbrain/core.py b/speechbrain/core.py
@@ -1295,9 +1295,16 @@ def evaluate_batch(self, batch, stage):
  -------
  detached loss
  """
-
- out = self.compute_forward(batch, stage=stage)
- loss = self.compute_objectives(out, batch, stage=stage)
+ amp = AMPConfig.from_name(self.precision)
+ if self.use_amp:
+ with torch.autocast(
+ dtype=amp.dtype, device_type=torch.device(self.device).type,
+ ):
+ out = self.compute_forward(batch, stage=stage)
+ loss = self.compute_objectives(out, batch, stage=stage)
+ else:
+ out = self.compute_forward(batch, stage=stage)
+ loss = self.compute_objectives(out, batch, stage=stage)
  return loss.detach().cpu()
 
  def _fit_train(self, train_set, epoch, enable):