Can I implement a custom loss function like this?

Dear Pytorch forum, I was wondering if you could help me with an implementation of a custom loss function. In particular, I want to perform some calculations on the model weights during my loss function - for example L1 penalty. My boundary condition is that my loss function is written as a nn.Module. Why? Because I wish to use Skorch for hyper parameter tuning, as Skorch only accepts nn.Module loss module as inputs (at least this is to the best of my understanding. I am aware, that if I wrote my NN in a Pytorch-typical loop, this problem would be much easier as I wouldn’t need to write a class for my loss, however, that way hyper parameter tuning is more cumbersome.

Assume I have a NN defined as follows:

import torch
import torch.nn as nn

class NeuralNet(nn.Module):
    trivial network, please do not read too much into it - its a matter of principle.
    Architecture: some trivial constant number of nodes for each HL, given number of nodes and HL.
    input_dim: dimension of input layer, hl: number of hidden layers, nodes: number of nodes per hl
    def __init__(self, input_dim, hl, nodes):
        self.input_dim = input_dim
        self.hl = hl
        self.nodes = nodes
        self.layers = nn.ModuleList()
        for h in range(hl):
            if h == 0:  # input layer
                self.layers.append(nn.Linear(input_dim, nodes))
                self.layers.append(nn.Linear(nodes, nodes))
        self.layers.append(nn.Linear(nodes, 1))
        self.model = nn.Sequential(*self.layers)

    def _get_weights():
        Gets model weights only and ignores biases
        note: w = [*layers.parameters()] would include biases I think
        w = []
        for name, param in self.layers.named_parameters():
            if 'weight' in name:
        return w

    def forward(self, x):
        Unusual, because it returns 2 elements. why? because I want the weights to be passed to my loss function.
        w = self._get_weights()
        x = self.model(x)
        return x, w

class RegMSE(nn.Module):
    Custom loss function.
    Example: perform L1 penalty on weights. Please ignore if this would make sense, it's a matter of principle.
    def __init__(self, l1):
        super(RegMSE, self).__init__()
        self.l1 = l1

    def forward(self, y_pred, target):
        y_pred now includes the weights also from the forward pass
        # unpack y_pred
        y_hat, weights = y_pred
        # calculate typical MSE loss
        loss = 1/(y_hat.shape[0])*torch.sum((target-y_hat) ** 2)
        l = torch.tensor([0], dtype=torch.float32)
        # L1 penalty
        if self.l1 is not None:
            for p in weights:
                l += torch.norm(p, 1)
            loss += self.l1 * l[0]
        return loss

Now in a training loop, we could have a situation where:

model = NeuralNet(input_dim=10, hl=2, nodes=10)
criterion = RegMSE(l1=1e-4)
ptimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)  

out = model(X)
loss = criterion(out, Y)

would this work to implement a L1 penalty? Meaning, would my weights be trained accordingly?

Because consequently, I could use it in skorch:

from skorch import NeuralNetRegressor

net = NeuralNetRegressor(
    module__input_dim=X.shape[1],  # Input shape
    module__nodes=10,  # Number of nodes in first hidden layer
    module__hl=5,  # Number of hidden layers
    criterion=RegMSE,  # Custom Loss Function
    criterion__l1=1e-4,  # L1 weight penalty
    batch_size=20, # Batch size
    max_epochs=100, # Epochs
    optimizer=torch.optim.Adam,  # Oprimizer
    optimizer__lr=1e-6,  # Learning Rate
    optimizer__weight_decay=0,  # No weight penalisation in optimizer
    iterator_train__drop_last=True,  # Avoid empty batch
    verbose=True  # Print learning

I want to use skorch to find an “optimal” l1 penalty using RandomizedSearchCV for example.

I really hope this makes sense, and I am grateful for any feedback! Perhaps this is not a performant solution, so if you can think of better ways of doing this, please let me know.