Fix to work in batch mode

[dapperdappy]

I'm not sure if this messes anything else up, but this is what I needed to run this script in the batch mode. It was complaining with this error:

RuntimeError: The size of tensor a (4) must match the size of tensor b (3) at non-singleton dimension 0