(arg="auto")
| 111 | |
| 112 | |
| 113 | def pick_device(arg="auto"): |
| 114 | if arg != "auto": |
| 115 | return torch.device(arg) |
| 116 | if torch.cuda.is_available(): |
| 117 | return torch.device("cuda") |
| 118 | if torch.backends.mps.is_available(): |
| 119 | return torch.device("mps") |
| 120 | return torch.device("cpu") |
| 121 | |
| 122 | |
| 123 | def quit_if_window_closed(env): |