Skip to content

Added MPS support to trainer.py#338

Open
leewinder wants to merge 2 commits intoace-step:mainfrom
leewinder:main
Open

Added MPS support to trainer.py#338
leewinder wants to merge 2 commits intoace-step:mainfrom
leewinder:main

Conversation

@leewinder
Copy link

Current trainer.py is hardcoded for CUDA support so this PR adds support for MPS and non-CUDA devices.

Also added in the --save_top_k srg as I kept running out of disk space on long training runs.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant