Update test_train_spmd_linear_model.py to show clear loss value

pytorch · Dec 11, 2024 · 00de506 · 00de506
1 parent 7dd2697
commit 00de506
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/spmd/test_train_spmd_linear_model.py b/test/spmd/test_train_spmd_linear_model.py
@@ -45,10 +45,10 @@ def __init__(self):
     super(SimpleLinear, self).__init__()
     self.fc1 = nn.Linear(FLAGS.input_dim, FLAGS.input_dim // 2)
     self.relu = nn.ReLU()
-    self.fc2 = nn.Linear(FLAGS.input_dim // 2, 1)
+    self.fc2 = nn.Linear(FLAGS.input_dim // 2, 3)
     # Add an additional 1x1 layer at the end to ensure the final layer
     # is not sharded.
-    self.fc3 = nn.Linear(1, 1)
+    self.fc3 = nn.Linear(3, 3)
 
   def forward(self, x):
     y = self.relu(self.fc1(x))