fix existing errors

pytorch · Sep 15, 2020 · 54b6c14 · 54b6c14
1 parent 9282009
commit 54b6c14
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/torchvision/models/googlenet.py b/torchvision/models/googlenet.py
@@ -1,5 +1,6 @@
 import warnings
 from collections import namedtuple
+from typing import Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -193,8 +194,7 @@ def _forward(self, x):
         return x, aux2, aux1
 
     @torch.jit.unused
-    def eager_outputs(self, x, aux2, aux1):
-        # type: (Tensor, Optional[Tensor], Optional[Tensor]) -> GoogLeNetOutputs
+    def eager_outputs(self, x: Tensor, aux2: Tensor, aux1: Optional[Tensor]) -> Union[GoogLeNetOutputs, Tensor]:
         if self.training and self.aux_logits:
             return _GoogLeNetOutputs(x, aux2, aux1)
         else:

diff --git a/torchvision/models/inception.py b/torchvision/models/inception.py
@@ -1,5 +1,6 @@
 from collections import namedtuple
 import warnings
+from typing import Union
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -188,8 +189,7 @@ def _forward(self, x):
         return x, aux
 
     @torch.jit.unused
-    def eager_outputs(self, x, aux):
-        # type: (Tensor, Optional[Tensor]) -> InceptionOutputs
+    def eager_outputs(self, x: torch.Tensor, aux: Optional[Tensor]) -> Union[InceptionOutputs, torch.Tensor]:
         if self.training and self.aux_logits:
             return InceptionOutputs(x, aux)
         else: