If you are talking about this training, I think you can just change it here:

  
      github.com
  
  
    pytorch/examples/blob/4581968193699de14b56527296262dd76ab43557/word_language_model/main.py#L104

###############################################################################
# Build the model
###############################################################################


ntokens = len(corpus.dictionary)
if args.model == 'Transformer':
    model = model.TransformerModel(ntokens, args.emsize, args.nhead, args.nhid, args.nlayers, args.dropout).to(device)
else:
    model = model.RNNModel(args.model, ntokens, args.emsize, args.nhid, args.nlayers, args.dropout, args.tied).to(device)


criterion = nn.CrossEntropyLoss()


###############################################################################
# Training code
###############################################################################


def repackage_hidden(h):
    """Wraps hidden states in new Tensors, to detach them from their history."""


    if isinstance(h, torch.Tensor):
        return h.detach()