Model fails to train and gives error when num_categories=1

I have a simple code from the demo, however no training happens. This is for segmentation. The model output has pixels of same value

```python
def load_model():
weights_manager = satlaspretrain_models.Weights()
model = weights_manager.get_pretrained_model(model_identifier="Aerial_SwinB_SI" , fpn=True , num_categories=2 , head=satlaspretrain_models.Head.SEGMENT).to(device)
return model

model = load_model()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)

val_step = 1 # evaluate every val_step epochs
for epoch in range(num_epochs):
print("Starting Epoch...", epoch)
for data, target in train_dataloader:
data = data.to(device)
target = target.to(device)

output, loss = model(data, target)
print("Train Loss = ", loss)
loss.backward()
optimizer.step()
# optimizer.zero_grad()

```

Also when `num_categories=1` I get an error
`model = weights_manager.get_pretrained_model(model_identifier="Aerial_SwinB_SI" , fpn=True , num_categories=1 , head=satlaspretrain_models.Head.SEGMENT).to(device)`
Error
` return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
IndexError: Target 1 is out of bounds.
`

I believe both these errors are related to how we define the loss function but I'm not sure. Can you please help me with the correct usage.

Update:
Changed the model to only use backbone, training occurs but is not very effective
```
class CustomModel(nn.Module):
def __init__(self, backbone, num_classes=1):
super(CustomModel, self).__init__()

# The backbone model (your pretrained model)
self.backbone = backbone

# Adjust upsampling layers to ensure output sizes match
self.upsample_1 = nn.ConvTranspose2d(128, 128, kernel_size=4, stride=4, padding=0)
self.upsample_2 = nn.ConvTranspose2d(128, 128, kernel_size=8, stride=8, padding=0)
self.upsample_3 = nn.ConvTranspose2d(128, 128, kernel_size=16, stride=16, padding=0)
self.upsample_4 = nn.ConvTranspose2d(128, 128, kernel_size=32, stride=32, padding=0)

# self.final_conv = nn.Conv2d(640, 1, kernel_size=1)
self.final_conv = nn.Sequential(
nn.Conv2d(640, 256, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(256, 64, kernel_size=3, padding=1),
nn.ReLU(),
nn.Conv2d(64, 1, kernel_size=1)
)

def forward(self, x):
# Get the feature maps from the backbone
out = self.backbone(x)

# Upsample the feature maps
out[1] = self.upsample_1(out[1])
out[2] = self.upsample_2(out[2])
out[3] = self.upsample_3(out[3])
out[4] = self.upsample_4(out[4])

# Concatenate the upsampled feature maps
combined = torch.cat([out[0], out[1], out[2], out[3], out[4]], dim=1)

# Apply a final convolution layer
output = self.final_conv(combined)

return output

# Instantiate the model with the backbone and attach the custom head
# backbone = weights_manager.get_pretrained_model(
# model_identifier="Aerial_SwinB_SI",
# fpn=True,
# num_categories=1,
# device='cpu',
# head=None
# ).to(device)

# # Create the custom model with added layers
# model = CustomModel(backbone).to(device)
```