remove assert statement as we added the check in testing

pytorch · Dec 4, 2024 · 2675b4c · 2675b4c
1 parent cd8be64
commit 2675b4c
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/torchtune/modules/early_exit_loss.py b/torchtune/modules/early_exit_loss.py
@@ -135,7 +135,6 @@ def layer_ids_to_loss_scales(
     loss_scales = loss_scales * torch.where(layer_ids < n_layers - 1, e_scale, 1.0)
     # normalize loss scales to ensure that their sum is 1.0
     loss_scales = loss_scales / torch.sum(loss_scales)
-    assert torch.isclose(torch.sum(loss_scales), torch.Tensor([1.0]).to(loss_scales))
 
     return loss_scales