gulabpatel/First-Order-Motion

1# -*- coding: utf-8 -*-2# File   : unittest.py3# Author : Jiayuan Mao4# Email  : maojiayuan@gmail.com5# Date   : 27/01/20186# 7# This file is part of Synchronized-BatchNorm-PyTorch.8# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch9# Distributed under MIT License.10 11import unittest12 13import numpy as np14from torch.autograd import Variable15 16 17def as_numpy(v):18    if isinstance(v, Variable):19        v = v.data20    return v.cpu().numpy()21 22 23class TorchTestCase(unittest.TestCase):24    def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3):25        npa, npb = as_numpy(a), as_numpy(b)26        self.assertTrue(27                np.allclose(npa, npb, atol=atol),28                'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max())29        )30