torch显卡或cpu转换自动推断函数

import torch
import numpy as np

def dynamic_planning_hardware(arr):
    decision_making = torch.cuda.is_available()
    if isinstance(arr,(list,np.ndarray)):
        if decision_making:
            return torch.tensor(arr).cuda()
        else:
            return torch.tensor(arr).cpu()
    elif arr.type() == 'torch.LongTensor' or arr.type() == 'torch.IntTensor' or arr.type() == 'torch.FloatTensor':
        if decision_making:
            return arr.cuda()
        else:
            return arr.cpu()

猜你喜欢

转载自blog.csdn.net/weixin_43069769/article/details/108047721
今日推荐