Solved by ptrblck in post #2 
              
              
                Sure, as long as you use PyTorch operations, you should be fine. 
Here is a dummy implementation of nn.MSELoss using the mean: 
def my_loss(output, target):
    loss = torch.mean((output - target)**2)
    return loss

model = nn.Linear(2, 2)
x = torch.randn(1, 2)
target = torch.randn(1, 2)
output = …