From ff2575b393e207f1cf4dc728cbd22009da6f9db5 Mon Sep 17 00:00:00 2001
From: gent <jw02425@surrey.ac.uk>
Date: Fri, 12 Jan 2024 17:25:53 +0000
Subject: [PATCH] fix compile for eval_cls.py
---
vitookit/evaluation/eval_cls.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/vitookit/evaluation/eval_cls.py b/vitookit/evaluation/eval_cls.py
index 37d5d68..b5cf20e 100644
--- a/vitookit/evaluation/eval_cls.py
+++ b/vitookit/evaluation/eval_cls.py
@@ -335,7 +335,7 @@ def main(args):
if args.pretrained_weights:
load_pretrained_weights(model, args.pretrained_weights, checkpoint_key=args.checkpoint_key, prefix=args.prefix)
if args.compile:
- return torch.compile(model)
+ model = torch.compile(model)
trunc_normal_(model.head.weight, std=2e-5)
model.to(device)
--
GitLab