use union for py 3.9 compatibilty

michael-diggin · pytorchmergebot · commit 6d1b76f8897a · 2024-11-15T23:10:15.000Z
diff --git a/torch/nn/functional.py b/torch/nn/functional.py
@@ -3230,7 +3230,7 @@ def poisson_nll_loss(
 def gaussian_nll_loss(
     input: Tensor,
     target: Tensor,
-    var: Tensor | float,
+    var: Union[Tensor, float],
     full: bool = False,
     eps: float = 1e-6,
     reduction: str = "mean",
diff --git a/torch/nn/modules/loss.py b/torch/nn/modules/loss.py
@@ -1,5 +1,5 @@
 # mypy: allow-untyped-defs
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 from typing_extensions import deprecated
 
 from torch import Tensor
@@ -438,7 +438,9 @@ def __init__(
         self.full = full
         self.eps = eps
 
-    def forward(self, input: Tensor, target: Tensor, var: Tensor | float) -> Tensor:
+    def forward(
+        self, input: Tensor, target: Tensor, var: Union[Tensor, float]
+    ) -> Tensor:
         return F.gaussian_nll_loss(
             input, target, var, full=self.full, eps=self.eps, reduction=self.reduction
         )
