diff --git a/vitookit/evaluation/eval_linear_ffcv.py b/vitookit/evaluation/eval_linear_ffcv.py index b37c0792811624006d8bd680178117c8213feff7..26cabe0db90025c720c8f88d23efc12feecc9e9b 100644 --- a/vitookit/evaluation/eval_linear_ffcv.py +++ b/vitookit/evaluation/eval_linear_ffcv.py @@ -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