search

angelolab / Nimbus-Inference

MIT License
12 stars 3 forks source link

MPS backend support and Dependency fix #5

Closed srivarra closed 7 months ago

srivarra commented 7 months ago

If you haven't already, please read through our contributing guidelines before opening your PR

What is the purpose of this PR?

Adds the ark-analysis package as a dependency to the project in order to run the notebook.

Adds PyTorch MPS support to Nimbus.

How did you implement your changes

Adds a new parameter in the Nimbus nn.Module derived class called device. Defaults to "auto" which will use the GPU if available (either Nvidia or Apple Silicon), otherwise runs on CPU.

Remaining issues

N/A

review-notebook-app[bot] commented 7 months ago

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB