Fix gradient accumulation

[rockerBOO]

Fixes #6 #2

• Tests added to showcase that it is working.
• Added robustness tests to infer that it is working as before.
• Test that tensors are on the same device (when we move tensors around or offload before starting the optimizer)
• Use autocast to make sure we handle mixed-precision appropriately and cast to float32 when appropriate

Use requirements-test.txt for testing requirements.

I think these changes do fix the issues but a performance overview will be important to consider. The tests should give a good grounding to make sure things continue to work.