I went via the second way. For example, here is my code for Linear layer:
import torch.nn.functional as F

class LinearOp:
    def __init__(self, weights, bias=None):
        self.weights = weights
        self.bias = bias

    def __call__(self, X):
        return F.linear(X, self.weights, self.bias)

It’s not handy to write such a class for each layer (convolutional, batchnorm, etc)