Skip to content
Snippets Groups Projects
Commit f07c128d authored by gent's avatar gent
Browse files

Import torch module and fix seed for reproducibility

parent ef8c2075
No related branches found
No related tags found
No related merge requests found
......@@ -25,9 +25,7 @@ import wandb
from vitookit.models.build_model import build_model
from vitookit.utils import misc
from vitookit.datasets import build_dataset
from torchvision.transforms import *
from vitookit.utils.helper import aug_parse, load_pretrained_weights, restart_from_checkpoint
from timm.models.layers import trunc_normal_
......@@ -128,6 +126,7 @@ def main(args):
print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
print("{}".format(args).replace(', ', ',\n'))
import torch
device = torch.device(args.device)
# fix the seed for reproducibility
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment