Thanks for the reply!
I will clarify a little. The output_logits in the code below are  torch.Size([128, 3]), a vector of 3 logit probabilities for every observation. My adjusted loss weights loss_mod is a 3*3 matrix.
I think I need the loss function to use this matrix along with the input and target. If I correctly understand how optimizing works, the direction in which the parameters are changed depends on the derivative of the loss function, right?
Looking at it more closely, I am wondering how loss.backward() and optimizer.step() are actually connected…
        loss_function = functional.cross_entropy
        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate,)
#
#
        for images_batch, labels_batch in train_data_loader:
            images_batch = images_batch.to(device)
            labels_batch = labels_batch.to(device)
            optimizer.zero_grad()
            output_logits = model(images_batch)
            loss_mod = np.array([[0, 1/3, 1], [2/3, 0, 2/3], [1, 1/3, 0]])
            loss = loss_function(input=output_logits, target=labels_batch)

            loss.backward()
            optimizer.step()


Greatly appreciated looking at this!
Cheers!