The commenting was for the model to learn what the types were so it could translate it to gather. Doing that one shot is surprising to me because it is somewhat hard to get pytorch tensors to work out to the correct shapes.