def METHOD_NAME(model_cls):
    # "model_cls" is a class, need to instantiate
    model = model_cls()
    # export model for visualization
    if 'NNI_OUTPUT_DIR' in os.environ:
        torch.onnx.export(model, (torch.randn(1, 1, 28, 28), ),
                          Path(os.environ['NNI_OUTPUT_DIR']) / 'model.onnx')
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
    transf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    train_loader = DataLoader(MNIST('data/mnist', download=True, transform=transf), batch_size=64, shuffle=True)
    test_loader = DataLoader(MNIST('data/mnist', download=True, train=False, transform=transf), batch_size=64)
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    model.to(device)
    for epoch in range(3):
        # train the model for one epoch
        train_epoch(model, device, train_loader, optimizer, epoch)
        # test the model for one epoch
        accuracy = test_epoch(model, device, test_loader)
        # call report intermediate result. Result can be float or dict
        nni.report_intermediate_result(accuracy)
    # report final test result
    nni.report_final_result(accuracy)