cumtcsys

I have encountered the same problem as you,after changed the third line "grid_size = inp_dim // stride" to "grid_size = prediction.size(2)" in function predict_transform,the problem fixed.