| import torch | |
| from spectral_world_models.models import build_model | |
| from spectral_world_models.metrics import count_parameters | |
| m=build_model('swm') | |
| print('torch:', torch.__version__) | |
| print('cuda:', torch.cuda.is_available()) | |
| print('parameters:', count_parameters(m)) | |