显式调用了cuda以后,pytorch的张量还是没有转移带GPU上的解决方法

    技术2022-07-11  98

    在使用GPU加速计算的时候,有的时候将张量显式调用了以后会发现其实还是在cpu上计算

    打个比方:

    class A: def __init__(self): super(...) a = torch.randn((1,2,3)).cuda() OR class B: def __init(self): a = torch.randn((1,2,3)) b = B() b.cuda()

     

    在显式调用了cuda以后还是报错在CPU上面计算

    问题就是变量类型上面了,这就是调用了cuda以后pytorch并没有将张量转移到GPU上面去进行计算,这个时候我们需要将Tensor转为nn.Parameter类型即可:

    class A: def __init__(): a = nn.Parameter(torch.randn((1,2,3))) a = A() a.cuda()

    这样就会转移到GPU上面去计算了

    Processed: 0.010, SQL: 9