Transformer.forward_train calls
tgt_mask = self.generate_square_subsequent_mask(tgt.shape[0], tgt.device)
but the helper only accepts sz, so the training branch of
Transformer.forward raises
TypeError: Transformer.generate_square_subsequent_mask() takes 2
positional arguments but 3 were given
The other two call sites (forward_test and forward_beam) pass one
argument and are unaffected.
Accept an optional device and forward it to the two tensor
constructors. It defaults to None, so the one-argument call sites keep
building the mask on the CPU exactly as before, and the call site that
already asks for a device now gets one.
