Update test_train_spmd_linear_model.py to show clear loss value (#8481)

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)
-    # Add an additional 1x1 layer at the end to ensure the final layer
+    self.fc2 = nn.Linear(FLAGS.input_dim // 2, 3)
+    # Add an additional 3x3 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))
@@ -121,6 +121,7 @@ def train_loop_fn(loader, epoch):
         optimizer.step()
       xm.mark_step()
       if step % 10 == 0:
+        assert loss != 0, "Loss should not 0 here"
         print(f"Epoch {epoch} step {step} loss {loss}")
 
   for epoch in range(FLAGS.num_epochs):