(data, yv_expect, running_mean, running_var)
| 33 | |
| 34 | @dist.launcher(n_gpus=2) |
| 35 | def worker(data, yv_expect, running_mean, running_var): |
| 36 | with amp.autocast(enabled=enable_amp): |
| 37 | rank = dist.get_rank() |
| 38 | bn = SyncBatchNorm(nr_chan, momentum=momentum, eps=eps) |
| 39 | for i in range(steps): |
| 40 | yv = bn(Tensor(data[rank][i])) |
| 41 | if enable_amp: |
| 42 | np.testing.assert_allclose( |
| 43 | yv.numpy(), yv_expect[rank], atol=5e-4, rtol=5e-4 |
| 44 | ) |
| 45 | else: |
| 46 | _assert_allclose(yv.numpy(), yv_expect[rank]) |
| 47 | _assert_allclose(bn.running_mean.numpy(), running_mean) |
| 48 | _assert_allclose(bn.running_var.numpy(), running_var) |
| 49 | |
| 50 | xv = [] |
| 51 | for i in range(steps): |
no test coverage detected