It seems that the computation of the survival layer in MCAT_Surv(link) is wrong, and logits = self.classifier(h).unsqueeze(0) should be logits = self.classifier(h). With the old version, supposing that the batch_size=6 and n_classes=4, the logits will be of size of (1,6,4), the hazards will be of size of (1,6,4), the Y_hat will be of size of (1,1,4), which certainly does not contain the Y_hat for the 6 samples of the batch. Besides, the S will means the cumulative production of the survival(i.e. 1-hazards) along the batch dimension, what does this mean? This S is of size of (1,6,4), then the len(S) in CoxSurvLoss(link) will be 1, which certainly is not the batch size as expected.

In the end, could you provide the reference of the equations for you to write this cox loss?