File size: 1,148 Bytes
fe64bad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
from argparse import ArgumentParser # noqa
def add_misc_options(parser):
group = parser.add_argument_group('Miscellaneous options')
group.add_argument("--expname", default="exps", help="general directory to this experiments, use it if you don't provide folder name")
group.add_argument("--folder", help="directory name to save models")
def add_cuda_options(parser):
group = parser.add_argument_group('Cuda options')
group.add_argument("--cuda", dest='cuda', action='store_true', help="if we want to try to use gpu")
group.add_argument('--cpu', dest='cuda', action='store_false', help="if we want to use cpu")
group.add_argument('--device', default=None, type=int, help="# of gpu device")
group.set_defaults(cuda=True)
def adding_cuda(parameters):
import torch
if parameters["cuda"] and torch.cuda.is_available():
if parameters.get('device') is not None:
parameters["device"] = torch.device(f"cuda:{parameters['device']}")
else:
parameters["device"] = torch.device("cuda")
else:
parameters["device"] = torch.device("cpu")
|