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))