Ah, perfect! Yep, this makes sense. Thanks very much for explaining it.
I also want to offer a revision to my previous post. Turns out you get very bad model performance if you do that training sub-loop like I was. The correct way to do teacher forcing is just to pass the targets shifted left one. That way, since the attention mask already restricts the model from cheating by looking ahead, each output token will have to be correctly predicted based on all the previous tokens in the input. Effectively, the same thing is accomplished as my training sub-loop, but without the massive inefficiency and bad performance.
# Model requires both "inputs" and "targets"
opt.zero_grad()
in_tgt = targets[:, :-1]
exp_tgt = targets[:, 1:]
tgt_padding_mask = tgt_padding_mask[:, :-1]
# Some code missing here, assume in_tgt gets converted to in_tgt_emb
out = model(inp_emb, in_tgt_emb, tgt_mask=tgt_mask, src_key_padding_mask=inp_padding_mask, tgt_key_padding_mask=tgt_padding_mask)
loss = criterion(out, exp_tgt)
loss.backward()
opt.step()
sch.step()