def create_model(ema=False):
model = models.TCDCNN()
model = model.cuda()
if ema:
for param in model.parameters():
param.detach_()
return model
def __init__(self, model, ema_model, alpha=0.999):
self.model = model
self.ema_model = ema_model
self.alpha = alpha
self.tmp_model = models.TCDCNN().cuda()
self.wd = 0.02 * args.lr
for param, ema_param in zip(self.model.parameters(), self.ema_model.parameters()):
ema_param.data.copy_(param.data)
