Train a Neural Network and using it as a loss function

autograd
9

I feel one way I could achieve this is by writing the code like this. Could you verify the effect of the deepcopy from copy? What happens when the model is loaded on the GPU?

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import SGD
import copy

class Model(nn.Module):
    def __init__(self, grayscale=False):
        super().__init__()

        self.inp1 = nn.Linear(10, 2)
        self.inp2 = nn.Linear(10,2)
        self.out = nn.Linear(4, 10)

    def forward(self, im1, im2):
        im1 = self.inp1(im1)
        im2 = self.inp2(im2)
        return self.out(torch.concat([im1, im2], axis=1))

model = Model()
optimizer = SGD(model.parameters(), lr=0.01, momentum=0.9)
loss_model = copy.deepcopy(model)
inp1 = torch.randn(1, 10)
inp2 = torch.randn(1, 10)
gt = torch.randn(1, 10)

pred = model(inp1, inp2)
loss_model = copy.deepcopy(model)
pred_new = loss_model(pred, inp1)
loss = F.mse_loss(inp1, pred_new)
loss.backward()
optimizer.step()

Also the memory footprint of the training is super high due to the copy of the model with all its parameters. Is their an elegant method to prevent the copy of the model to achieve the same results?