torch.jit保存，加载模型

  

   

    

     

    
    

     

      ori_backbone.eval()
     
    
   

    

     

    
    

     

      x = torch.randn(1,3,112,112).cuda()
     
    
   

    

     

    
    

     

      traced_cell = torch.jit.trace(ori_backbone, (x))
     
    
   

    

     

    
    

     

      torch.jit.save(traced_cell, os.path.join(MODEL_ROOT, "Epoch_{}_Time_{}_checkpoint.pth".format(epoch + 1, get_time())))
     
    
  
 

  

   

    

     

    
    

     

      model=torch.jit.load(spath[0], map_location=tnet.ctx)
     
    
   

    

     

    
    

     

      model.eval()
     
    
  
 

Traceback (most recent call last):
  File "E:/project/shibie/cavaface.pytorch/demo.py", line 14, in <module>
    model= torch.jit.load(a_path)#, map_location=torch.device("cpu"))
  File "D:\Anaconda3\lib\site-packages\torch\jit\__init__.py", line 228, in load
    cpp_module = torch._C.import_ir_module(cu, f, map_location, _extra_files)
RuntimeError: Cannot initialize CUDA without ATen_cuda library. PyTorch splits its backend into two shared libraries: a

文章来源: blog.csdn.net，作者：网奇，版权归原作者所有，如需转载，请联系作者。

原文链接：blog.csdn.net/jacke121/article/details/115539086