make func

Signed-off-by: Mayank Mishra <[email protected]>
IBM · Sep 12, 2024 · b2c19d7 · b2c19d7
1 parent e893721
commit b2c19d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dolomite_engine/hf_models/mixins/moe/base.py b/dolomite_engine/hf_models/mixins/moe/base.py
@@ -9,7 +9,7 @@
 
 @dataclass
 class MoeModelOutputWithPastAndAuxLoss(MoeModelOutputWithPast):
-    aux_loss: torch.Tensor = None
+    aux_loss: torch.Tensor | None = None
 
 
 class BaseMoEModelMixin(BaseModelMixin):