insightface 切换GPU训练报错解决

insightface代码：

https://github.com/deepinsight/insightface

在pytorch训练时，先用gpu1训练，再用gpu2训练，加载gpu1的预训练时报错，

解决方法：

先用cpu加载模型，再转换到新的gpu id上，

代码如下：

      self.weight: torch.Tensor = torch.load(self.weight_name,map_location='cpu')

      self.weight=self.weight.cuda(local_rank)
      self.weight_mom: torch.Tensor = torch.load(self.weight_mom_name,map_location='cpu')
      self.weight_mom=self.weight_mom.cuda(local_rank)

partial_fc.py 完整代码如下：

import logging
import os

import torch
import torch.distributed as dist
from torch.nn import Module
from torch.nn.functional import normalize, linear
from torch.nn.parameter import Parameter


class PartialFC(Module):
    """
    Author: {Xiang