(self)
| 250 | |
| 251 | @skip_if_no_cuda |
| 252 | def test_amp(self): |
| 253 | shape = (1, 3, 10, 8) |
| 254 | device = "cuda" |
| 255 | im, _ = self.get_im(shape, device=device) |
| 256 | conv = torch.nn.Conv2d(im.shape[1], 5, 3) |
| 257 | conv.to(device) |
| 258 | im_conv = conv(im) |
| 259 | with torch.autocast("cuda"): |
| 260 | im_conv2 = conv(im) |
| 261 | self.check(im_conv2, im_conv, ids=False, rtol=1e-2, atol=1e-2) |
| 262 | |
| 263 | def test_out(self): |
| 264 | """Test when `out` is given as an argument.""" |