Cox-STG model cannot predict properly

[shi-ang]

Your Cox-STG model is not compatible with some methods in the super class.

For example, if I train a Cox-STG model and want to get the risk prediction (for clinical usage or compute downstream survival curves with Breslow estimator). neither
model.predict(test_data['X'])
nor
model.predict(test_data)
works.

This problem can be reproduced by adding these two lines at the end of the python/examples/Cox-example.ipynb file.

I use the following temporary code to predict the risk score:

test_loader = model.get_dataloader(test_data['X'], {'E': test_data['E'], 'T': test_data['T']}, shuffle=None)
for fd in test_loader:
    with torch.no_grad():
        pred = model._model(fd)
risks = pred['logits']

but I still would be very appreciate if you can fix this problem