Request for additional section of code regarding loss functions

Hello! Thank you for building this super beneficial and simple repository.

We are currently testing and training and we would like to incorporate the following collision loss function as well and would really appreciate if you could provide the same.

```
# [Collision loss]
# if classifier is not None:
# real_labels = torch.ones(joint.size(0), dtype=torch.long).to(joint.device)
# # Discriminator's output for generated data
# safe_logits = classifier(joint)
# criterion = nn.CrossEntropyLoss()
# # Generator loss is the cross-entropy loss between the fake outputs and the label 1 (real)
# collision_loss = criterion(safe_logits, real_labels)

# collision Loss integration pending.
collision_loss = torch.tensor([0.0]).cuda()
```

Awaiting your reply!