From f00a5b84b37e5e2935d69a60044e30cd1b46e92f Mon Sep 17 00:00:00 2001 From: Toni Kukurin Date: Sat, 10 Dec 2022 14:23:55 +0000 Subject: [PATCH] train all batches! --- utils.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/utils.py b/utils.py index d638390..0244296 100644 --- a/utils.py +++ b/utils.py @@ -521,16 +521,16 @@ class CCS(object): x0_batch = x0[j*batch_size:(j+1)*batch_size] x1_batch = x1[j*batch_size:(j+1)*batch_size] - # probe - p0, p1 = self.probe(x0_batch), self.probe(x1_batch) + # probe + p0, p1 = self.probe(x0_batch), self.probe(x1_batch) - # get the corresponding loss - loss = self.get_loss(p0, p1) + # get the corresponding loss + loss = self.get_loss(p0, p1) - # update the parameters - optimizer.zero_grad() - loss.backward() - optimizer.step() + # update the parameters + optimizer.zero_grad() + loss.backward() + optimizer.step() return loss.detach().cpu().item()