Thanks for the reply @ptrblck. I have checked the size of the x i.e output at each stage in internal convnet for forward. It is as follows:
1: torch.Size([128, 64, 61, 61])
2: torch.Size([128, 238144])
3: torch.Size([128, 2])

as I am quite new to PyTorch I am not getting how can I convert it into 4D as per your suggestion? Shall I use unsqueez? but I guess it expands the dimension by 1D only.
Any suggestions or solutions from your side is welcome. It would be great help.
Thanks in advance.