Thank you very much, that was the solution for me. I searched for someone who had the same model, but unfortunately, I am having some problems there with torch.cuda.amp. Could you suggest some projects using different inputs in same the model?