import torch import thop from nasnet import * x = torch.rand(1, 3, 224, 224) # model = pdarts(init_channel=48, layers=14) model = pdarts_handwritten() model.train()#eval() flops, params = thop.profile(model, (x,), verbose=True) print(flops, params) assert flops == 566209280, flops # assert params < 50e6