def get_model(opt):
network = ProtoNetEmbedding().cuda()
if opt.cls_head == 'R2D2':
cls_head = R2D2Head().cuda()
elif opt.cls_head == 'SVM':
cls_head = SVMHead().cuda()
return (network, cls_head)
def get_model():
network = ProtoNetEmbedding().cuda()
cls_head = ClassificationHead(enable_scale=False).cuda()
return (network, cls_head)