Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Motivation and context / Related issue

Resolves #284, Resolves #287
Tests on GPU were failing

How has this been tested (if it applies)

Tested on Linux

Checklist

• The documentation is up-to-date with the changes I made.
• I have read the CONTRIBUTING document.
• All tests passed, and additional code has been covered with new tests.