Basically you could handle it like a 2-dimensional convolution with another “spatial” dimension.
I.e. the target should contain the class indices without a channel dimension in the shape [batch_size, d, h, w].
I’ve created a small dummy example using a simple model to segment a square in the volume:
torch.manual_seed(2809)

c, d, h, w = 3, 12, 12, 12
data = torch.randn(1, c, d, h, w)
target = torch.zeros(1, 12, 12, 12, dtype=torch.long)
target[0, 4:8, 4:8, 4:8] = 1

model = nn.Sequential(
    nn.Conv3d(c, 6, 3, 1, 1),
    nn.ReLU(),
    nn.Conv3d(6, 2, 3, 1, 1)
)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(800):
    optimizer.zero_grad()
    
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    
    print('Epoch {}, loss {}'.format(epoch, loss.item()))


pred = torch.argmax(output.detach(), 1).squeeze(0)
pred.unique()
rows, cols = 4, 3
fig, axarr = plt.subplots(rows, cols)
for idx in range(pred.size(0)):
    axarr[idx//cols, idx%cols].imshow(pred[idx])

fig, axarr = plt.subplots(rows, cols)
for idx in range(target.size(1)):
    axarr[idx//cols, idx%cols].imshow(target[0, idx])

Let me know, if that helps in creating your model.