Htut_Lynn_Aung:

 def forward(self, x, target_size):
        # assert (x.data.dim() == 4)



this is an error from torch.jit.script, you can annotate the forward like following:
 def forward(self, x, target_size):
        # type : (Tensor, List[int]) -> Tensor
        # assert (x.data.dim() == 4)

to make it scriptable.