(arg="auto")
| 207 | |
| 208 | |
| 209 | def pick_device(arg="auto"): |
| 210 | if arg != "auto": |
| 211 | return torch.device(arg) |
| 212 | if torch.cuda.is_available(): |
| 213 | return torch.device("cuda") |
| 214 | if torch.backends.mps.is_available(): |
| 215 | return torch.device("mps") |
| 216 | return torch.device("cpu") |
| 217 | |
| 218 | |
| 219 | def quit_if_window_closed(env): |