Thanks, @allenye0119!
I can see the idea now, but I don’t see how to use the modified forward() function. The training code at each epoch looks like this:
    optimizer.zero_grad()
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()

So I guess I have to modify Model() so that it takes an additional argument: output = model(data, data_type)? Thanks again for your help.