def test_shape(self, input_param, input, expected_shape):
result = AddCoordinateChannels(**input_param)(input)
self.assertEqual(type(result), type(input))
if isinstance(result, torch.Tensor):
self.assertEqual(result.device, input.device)
self.assertEqual(list(result.shape), list(expected_shape))
assert_allclose(input[0, ...], result[0, ...])
def test_channel_dim(self, input_param, input):
with self.assertRaises(ValueError):
AddCoordinateChannels(**input_param)(input)
def test_shape(self, input_param, input, expected_shape):
result = AddCoordinateChannels(**input_param)(input)
self.assertEqual(list(result.shape), list(expected_shape))
np.testing.assert_array_equal(input[0, ...], result[0, ...])