使用gpu跑模型 todevice()

gpu: todevice

1)model層面    

model = Net(in_feats=params.hidden_size, 
            n_hidden=params.hidden_size, 
            n_hidden_layers=1,
            vocab_size=len(type2id)).to(device)

2)數據層面
 

for epoch in range(pa.EPOCH_NUM):    
    for batch_g, batch_e, batch_l in trainloader:                
        batch_g.to(device)        
        batch_e = batch_e.to(device)
        batch_l = batch_l.to(device)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章