This is how I thought I’d do it. Can’t test it at the moment. I still have to somehow tell the old net that it doesnt need gradients, but I’ll figure that out.
        ce_loss = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(new_model.parameters(), lr=0.0001)
        # strength of additional loss
        strength = 0.001

        for inputs, targets in train_loader:
            optimizer.zero_grad()
            old_output = old_model(inputs)
            new_output = new_model(inputs)
            standard_loss = ce_loss(new_output, torch.max(targets, 1)[1])
            additional_loss = ce_loss(new_output, torch.max(old_output, 1)[1])
            total_loss = standard_loss + strength * (1/additional_loss)
            # or: total_loss = standard_loss - strength * additional_loss

            total_loss.backward()
            optimizer.step()