pytorch cuda上tensor的定义 以及 减少cpu操作的方法

首页 / 文章分类 / 正文

cuda上tensor的定义

a = torch.ones(1000,1000,3).cuda() 

某一gpu上定义

cuda1 = torch.device('cuda:1') b = torch.randn((1000,1000,1000),device=cuda1) 

删除某一变量

del a 

在cpu定义tensor然后转到gpu

torch.zeros().cuda() 

直接在gpu上定义,这样就减少了cpu的损耗

torch.cuda.FloatTensor(batch_size, self.hidden_dim, self.height, self.width).fill_(0)