Update test_train_spmd_linear_model.py

test/spmd/test_train_spmd_linear_model.py

@@ -46,7 +46,7 @@ def __init__(self):
     self.fc1 = nn.Linear(FLAGS.input_dim, FLAGS.input_dim // 2)
     self.relu = nn.ReLU()
     self.fc2 = nn.Linear(FLAGS.input_dim // 2, 3)
-    # Add an additional 1x1 layer at the end to ensure the final layer
+    # Add an additional 3x3 layer at the end to ensure the final layer
     # is not sharded.
     self.fc3 = nn.Linear(3, 3)