The GeneratorExp is not supported in torch.jit:
if any(img.size != imgs[0].size for img in imgs)

To script the model you could wrap this code path with if not torch.jit.is_scripting() to ignore it.
Of course you would have to make sure to pass the right input before calling the model.