wassname/discovering_latent_knowledge
1
0
Fork 0
mirror of https://github.com/wassname/discovering_latent_knowledge.git synced 2026-09-09 11:21:22 +08:00
Code Issues Packages Projects Releases Wiki Activity

train all batches!

Browse Source
This commit is contained in:
Toni Kukurin
2022-12-10 14:23:55 +00:00
parent 1fb40eefbd
commit f00a5b84b3
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils.py
+8 -8
View File
@@ -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()