libtorch freeze error

[KongCang]

model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
    cfg, is_train=False
)

if cfg.TEST.MODEL_FILE:
    logger.info('=> loading model from {}'.format(cfg.TEST.MODEL_FILE))
    model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
else:
    model_state_file = os.path.join(
        final_output_dir, 'final_state.pth'
    )
    logger.info('=> loading model from {}'.format(model_state_file))
    model.load_state_dict(torch.load(model_state_file))

model = torch.nn.DataParallel(model, device_ids=cfg.GPUS).cuda()
input = torch.rand((1, 3, 384, 288)).cuda() 
outputs =model(input)#success
traced_script_module = torch.jit.trace(model, input)#*** AssertionError other models can freeze like this

[KongCang]

After remove the None layers, freeze is ok.