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