From 49077267c6bf54440058d0fdd7fc48b6bdfbc03e Mon Sep 17 00:00:00 2001 From: Rahul Mangalampalli Date: Thu, 9 Dec 2021 18:04:18 +0530 Subject: [PATCH] Changed code from line 178 to 181 Added the snipped of code ``classifier[1].out_features``` to support Mobilenet classifiers in the code while training. --- main_dino.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/main_dino.py b/main_dino.py index cade9873d..bd99c1bf6 100644 --- a/main_dino.py +++ b/main_dino.py @@ -175,7 +175,10 @@ def train_dino(args): elif args.arch in torchvision_models.__dict__.keys(): student = torchvision_models.__dict__[args.arch]() teacher = torchvision_models.__dict__[args.arch]() - embed_dim = student.fc.weight.shape[1] + if args.arch.find("mobile") == -1: + embed_dim = student.fc.weight.shape[1] + else: + embed_dim = student.classifier[1].out_features else: print(f"Unknow architecture: {args.arch}")