PyTorch custom loss function

Question

How should a custom loss function be implemented ? Using below code is causing error :

import torch
import torch.nn as nn
import torchvision
import torchvision.         


        

Answer 1

If you use torch functions you should be fine

import torch 

def my_custom_loss(output, target):
    loss = torch.mean((output-target*2)**3)
    return loss

# Forward pass to the Network
# then, 
loss.backward()