LR Scheduler fixes

[stephenroller]

Patch description
Loading a model trained with cosine or linear LR schedules was throwing an exception on PyTorch 1.6. This is because pytorch 1.6 broke our method for forwarding an LR schedule with step(num_epochs).

List of changes:

• Migrates us to the new sanctioned LR loading
• Works on both pytorch 1.4 and 1.6 (didn't bother testing 1.5.1)
• Adds a new test_agents/unigram model, which is a torch generator agent that always outputs fixed distributions. This is useful to have a TA that has extremely few parameters (tens or hundreds) and can be trained extremely fast. (Also added its own CI for this)
• Add end2end tests for select LR schedulers
• Add resume LR checks for select LR schedulers.

Testing steps

• Ran select unit tests in both 1.4 and 1.6.
• New CI

[github-actions]

Your PR contains a change to a task. Please paste the results of the following command into a comment:

python tests/datatests/test_new_tasks.py

[klshuster]

thanks for adding a bunch of tests too