explosion/spaCy
1
1
Fork 0
mirror of https://github.com/explosion/spaCy.git synced 2024-11-14 21:57:15 +03:00
Code Issues Packages Projects Releases Wiki Activity

Enable GPU in pytorch n use_gpu functon

Browse Source
This commit is contained in:
Matthw Honnibal 2019-10-06 19:24:21 +02:00
parent 9dbaea1ab4
commit 63ff233ba2
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
spacy/util.py
View File

@ -707,6 +707,11 @@ def use_gpu(gpu_id):
device.use()
Model.ops = CupyOps()
Model.Ops = CupyOps
try:
import torch
torch.set_default_tensor_type("torch.cuda.FloatTensor")
except ImportError:
pass
return device