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
[X] 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)
Summary
This PR adds support for declaring number of clasees for Image Classification Pytorch models.
Details
Checklist
test-inputs.yml
, which specifies a few input values to test the code template (the test is then automatically run by pytest)