(self)
| 23 | |
| 24 | class TestPadMode(unittest.TestCase): |
| 25 | def test_pad(self): |
| 26 | expected_shapes = {3: (1, 15, 10), 4: (1, 10, 6, 7)} |
| 27 | devices = ("cuda:0", "cpu") if torch.cuda.is_available() else ("cpu",) |
| 28 | shapes = ((1, 10, 10), (1, 5, 6, 7)) |
| 29 | types = (float, int, np.uint8, np.int16, np.float32, bool) |
| 30 | modes = list(PytorchPadMode) + list(NumpyPadMode) |
| 31 | |
| 32 | for params in dict_product(t=types, d=devices, s=shapes, m=modes): |
| 33 | t = params["t"] |
| 34 | d = params["d"] |
| 35 | s = params["s"] |
| 36 | m = params["m"] |
| 37 | a = torch.rand(s) |
| 38 | to_pad = [(0, 0), (2, 3)] if len(s) == 3 else [(0, 0), (2, 3), (0, 0), (0, 0)] |
| 39 | out = Pad(to_pad=to_pad, mode=m)(CastToType(dtype=t)(a).to(d)) |
| 40 | self.assertEqual(out.shape, expected_shapes[len(s)]) |
| 41 | |
| 42 | |
| 43 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected