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()
torch显卡或cpu转换自动推断函数
猜你喜欢
转载自blog.csdn.net/weixin_43069769/article/details/108047721
今日推荐
周排行