self.CUDA = torch.cuda.is_available()
self.device = torch.device("cuda:0" if self.CUDA else "cpu")
model = PerfGuard(features1.shape[2],config.embd_dim, config.tensor_dim,config.dropout).cuda(config.device)
A1 = torch.from_numpy(preprocess_adj(A1)).float().cuda(config.device)
X1 = self.dropout(torch.from_numpy(X1).cuda(config.device))
将模型和数据都放在gpu上
最新推荐文章于 2023-12-02 10:17:42 发布