Method tearDown

tests/networks/test_convert_to_trt.py:46–49

(self)

        self.gpu_device = torch.cuda.current_device()

    def tearDown(self):
        current_device = torch.cuda.current_device()
        if current_device != self.gpu_device:
            torch.cuda.set_device(self.gpu_device)

    @parameterized.expand([TEST_CASE_1, TEST_CASE_2])
    @unittest.skipUnless(has_torchtrt and has_tensorrt, "Torch-TensorRT is required for convert!")