Commit 39f2da0

committed

test backend

1 parent c2b9506 commit 39f2da0Copy full SHA for 39f2da0

File tree

+11

-6

lines changed

+11

-6

lines changed

Lines changed: 8 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -80,9 +80,14 @@ def reset_parameters(self):`
`80`	`80`
`81`	`81`	`def __top_k__(self, x_s, x_t):`
`82`	`82`	`r"""Memory-efficient top-k correspondence computation."""`
`83`		`- x_s, x_t = LazyTensor(x_s.unsqueeze(-2)), LazyTensor(x_t.unsqueeze(-3))`
`84`		`- S_ij = (-x_s * x_t).sum(dim=-1)`
`85`		`- return S_ij.argKmin(self.k, dim=2, backend=self.backend)`
	`83`	`+ x_s, x_t = x_s.unsqueeze(-2), x_t.unsqueeze(-3)`
	`84`	`+ if self.backend != 'test': # pragma: no cover`
	`85`	`+ x_s, x_t = LazyTensor(x_s), LazyTensor(x_t)`
	`86`	`+ S_ij = (-x_s * x_t).sum(dim=-1)`
	`87`	`+ return S_ij.argKmin(self.k, dim=2, backend=self.backend)`
	`88`	`+ else:`
	`89`	`+ S_ij = (x_s * x_t).sum(dim=-1)`
	`90`	`+ return S_ij.topk(self.k, dim=2)[1]`
`86`	`91`
`87`	`92`	`def __include_gt__(self, S_idx, s_mask, y):`
`88`	`93`	r"""Includes the ground-truth values in :obj:`y` to the index tensor

Lines changed: 3 additions & 3 deletions

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ def test_dgmc_repr():`
`29`	`29`	`def test_dgmc_on_single_graphs():`
`30`	`30`	`set_seed()`
`31`	`31`	`model = DGMC(psi_1, psi_2, num_steps=1)`
`32`		`- model.backend = 'CPU'`
	`32`	`+ model.backend = 'test'`
`33`	`33`	`x, e = data.x, data.edge_index`
`34`	`34`	`y = torch.arange(data.num_nodes)`
`35`	`35`	`y = torch.stack([y, y], dim=0)`
`@@ -68,7 +68,7 @@ def test_dgmc_on_single_graphs():`
`68`	`68`	`def test_dgmc_on_multiple_graphs():`
`69`	`69`	`set_seed()`
`70`	`70`	`model = DGMC(psi_1, psi_2, num_steps=1)`
`71`		`- model.backend = 'CPU'`
	`71`	`+ model.backend = 'test'`
`72`	`72`
`73`	`73`	`batch = Batch.from_data_list([data, data])`
`74`	`74`	`x, e, b = batch.x, batch.edge_index, batch.batch`
`@@ -88,7 +88,7 @@ def test_dgmc_on_multiple_graphs():`
`88`	`88`
`89`	`89`	`def test_dgmc_include_gt():`
`90`	`90`	`model = DGMC(psi_1, psi_2, num_steps=1)`
`91`		`- model.backend = 'CPU'`
	`91`	`+ model.backend = 'test'`
`92`	`92`
`93`	`93`	`S_idx = torch.tensor([[[0, 1], [1, 2]], [[1, 2], [0, 1]]])`
`94`	`94`	`s_mask = torch.tensor([[True, False], [True, True]])`

Comments

(0)