| 30 | |
| 31 | |
| 32 | class SameScaleSampler(Sampler): |
| 33 | def __init__(self, *args, **kwargs): |
| 34 | super(SameScaleSampler, self).__init__(*args, **kwargs) |
| 35 | |
| 36 | def sample(self, anchor, sample, *args, **kwargs): |
| 37 | assert anchor.size(0) == sample.size(0) |
| 38 | num_nodes = anchor.size(0) |
| 39 | device = anchor.device |
| 40 | pos_mask = torch.eye(num_nodes, dtype=torch.float32, device=device) |
| 41 | neg_mask = 1. - pos_mask |
| 42 | return anchor, sample, pos_mask, neg_mask |
| 43 | |
| 44 | |
| 45 | class CrossScaleSampler(Sampler): |