King Yiu Suen

I don't have much experience with onxx, but it seems like you just need to add a forward method in LitResNetTransformer: ```python class LitResNetTransformer(LightningModule): def forward(self, x: Tensor): return self.model.predict(x)...

Based on the error message, it seems like the problem comes from L157 of models/resnet_transformer.py, since torch.all is used there. My guess is torch.jit.trace does not like conditionals on data...