moco_q_k.py

# zero out grads
optimizer.zero_grad()

# retrieve xq and xk the two image batches
xq = sample_batched['image1']
xk = sample_batched['image2']

# move them to the device
xq = xq.to(device)
xk = xk.to(device)

# get their outputs
q = resnetq(xq)
k = resnetk(xk)
k = k.detach()

# normalize the ouptuts, make them unit vectors
q = torch.div(q,torch.norm(q,dim=1).reshape(-1,1))
k = torch.div(k,torch.norm(k,dim=1).reshape(-1,1))