I think even if you don’t use PyTorch layers you should use Parameters for your learnable weights/params. This will probably make it easier for you.
Using a subclass of nn.Module will ensure that all the backend work for computing gradients will work out of the box, and you can use your functions (or let’s call it a module, then) like another module of pytorch layers…