| import torch | |
| def count_parameters(model): | |
| return sum( | |
| parameter.numel() | |
| for parameter in model.parameters() | |
| if parameter.requires_grad | |
| ) | |
| def get_device(): | |
| if torch.cuda.is_available(): | |
| return torch.device("cuda") | |
| if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): | |
| return torch.device("mps") | |
| return torch.device("cpu") |