loss

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Nov 4, 2024 · 4de6193 · 4de6193
1 parent 9545480
commit 4de6193
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dolomite_engine/hf_models/loss.py b/dolomite_engine/hf_models/loss.py
@@ -18,7 +18,7 @@ def get_autoregressive_language_modeling_loss(
     use_padding_free_transformer: bool = False,
     reduction: str = "mean",
     tensor_parallel_word_embeddings: bool = False,
-) -> DTensor:
+) -> torch.Tensor | DTensor:
     if use_padding_free_transformer:
         assert cu_seqlens is not None