Allow custom model when given by the user.

[leoniloris]

• The user can provide a custom pytorch model.
• Avoid repeating code on a few places.

[maxiimilian]

This looks very reasonable. Any plans to merge this request, @SJ001 ?