@JesseLivezey this can be done in a multiple way:

You would like to reuse create_supervised_trainer, you can create a custom function with a signature (y_pred, y)

model = ...
loss_1 = nn.CrossEntropy()
model_loss = ...

def my_custom_loss(y_pred, y):
    return loss_1(y_pred, y) + model_loss(model)

trainer = create_supervised_trainer(model, optimizer, my_custom_loss)


You would like to reuse create_supervised_trainer, you can subclass a loss function as

class ComposedLoss(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.l1 = nn.CrossEntropy()
        self.model_loss = ... 
        self.model = model

    def forward(self, y_pred, y):
         return self.l1(y_pred, y) + self.model_loss(self.model)
        


You can directly create a trainer with you own update function

def update_fn(engine, batch):
        model.train()
        optimizer.zero_grad()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        if model_loss is not None:
            loss = loss + model_loss(model)
        loss.backward()
        optimizer.step()
        return output_transform(x, y, y_pred, loss)

trainer = Engine(update_fn)

HTH