Yes you can do that. See more details in the doc 154.
Be careful:

You should not instantiate an instance of the Function. You should do loss = MyLoss.apply(output, label). If you prefer to instantiate an instance, you can simply create a new module like:

class MyLossMod(nn.Module):
    def forward(self, input, target):
        return MyLoss.apply(input, target)

# Then you can use any of the two
loss = MyLoss.apply(input, target)
# or
mylossmod = MyLossMod()
loss = mylossmod(input, target)


Your arguments for forward should be ctx not self but I guess that’s just a typo