The STN cropping in annotated_prtr.ipynb only works for batch_size=1

[EckoTan0804]

Hello,

I've tried the code of STN cropping in annotated_prtr.ipynb. For single image (i.e. batch_size=1) it works well. But for batch_size > 1 (e.g. 2) an error occurs:

(In my case, person_per_image is set to be 100, therefore the first dimension is 2 * 100 = 200)

It seems that this error is caused by expand() since expand() can only expand the axis with size 1 (according to the discussion: https://discuss.pytorch.org/t/runtimeerror-the-expanded-size-of-the-tensor-must-match-the-existing-size/105523).