Well,sorry, I just asnwer my own question - after going round the web and searching - it seems that simply changing the model mode from .train() to .eval() like learner.model.eval() fixes the issue.
1 Like