Change Dimension of Softmax from 0 to 1 in modules from week 1 to 4

[Aarish Alam]

The output logits are of shape (1,2) therefore the dimension while calling the softmax function should be 1. These line from week 1 to 4 needs to be changed in this file