Skip to content

GitLab

Commit b44ea8db authored by wanglch's avatar wanglch
Browse files

Update umt5_summary.py

parent 109633be
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
umt5_summary.py
View file @ b44ea8db
...@@ -4,7 +4,7 @@ from transformers import AutoModelForSeq2SeqLM ...@@ -4,7 +4,7 @@ from transformers import AutoModelForSeq2SeqLM
from torch import nn from torch import nn
import os import os
os.environ["HIP_VISIBLE_DEVICES"] = "4,5" os.environ["HIP_VISIBLE_DEVICES"] = "0"
device = 'cuda' if torch.cuda.is_available() else 'cpu' device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'Using {device} device') print(f'Using {device} device')
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment