main.py

from windformer import WindFormerDist
import torch


if __name__ == '__main__':
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    net_path = './windformer.pkl'

    model = WindFormerDist(fusion=False).to(device)
    model.load_state_dict(torch.load(net_path, map_location=device)['model_weights'])
    model.eval()