Update test_train_spmd_linear_model.py to show clear loss value

[ManfeiBai]

linear model return one dimension value for each input will cause loss not clear, we notice all 0 in loss regression;

so increase return value dimension for used torch.nn.CrossEntropyLoss() to figure out

local passed log: https://gist.github.com/ManfeiBai/df6362c9a76890847127374e60fd6846

[tengyifei]

Can we add an assertion on line 124 that checks loss is not 0?

[ManfeiBai]

Can we add an assertion on line 124 that checks loss is not 0?

Thanks, nice catch, updated with assertion