Type annotate dataloader2 train loop (#664)

Summary: Pull Request resolved: #664 * Type annotate dataloader2 train loop Reviewed By: Miiira Differential Revision: D37971478 fbshipit-source-id: 67cdd33fce70b43da4f9782f277510d1d7cfc7a1
pytorch · Jul 20, 2022 · b6ecce0 · b6ecce0
1 parent c2223fd
commit b6ecce0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/dataloader2/train_loop.py b/examples/dataloader2/train_loop.py
@@ -11,7 +11,7 @@
 
 
 class ToyModel(torch.nn.Module):
-    def __init__(self):
+    def __init__(self) -> None:
         """
         In the model constructor, we instantiate four parameters and use them
         as member parameters.
@@ -22,7 +22,7 @@ def __init__(self):
         self.c = torch.nn.Parameter(torch.randn(()))
         self.d = torch.nn.Parameter(torch.randn(()))
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         Simple model forward function
         """