Arbitrary class training in Pytorch #5

[murthy95]

Summary

This PR adds support for declaring number of clasees for Image Classification Pytorch models.

Details

• Adds a parameter in the sidebar below the "Use pre-trained model" checkbox
• Default value is 1000 (what torchvision uses)
• If pre-trained is True, imagenet weights are loaded and last fc layer is redefined

Checklist

• all tests are passing (see README.md on how to run tests)
• if you created a new template: it contains a file test-inputs.yml, which specifies a few input values to test the code template (the test is then automatically run by pytest)
• you formatted all code with black
• you checked all new functionality live, i.e. in the running web app
• any generated code is formatted nicely, both in .py and in .ipynb ("nicely" = comparable to the existing templates)
• you added comments in your code that explain what it does
• the PR explains in detail what's new

[jrieke]

Looks good! Merging for now, will do a few minor changes and let you know in the issue when it's deployed :)