From 73a14d7266f6246a573875e04ecd5306c894062e Mon Sep 17 00:00:00 2001
From: Jeremy Reizenstein <reizenstein@fb.com>
Date: Tue, 26 Oct 2021 14:32:34 -0700
Subject: [PATCH] dataparallel fix

Summary: Attempt to overcome flaky test

Reviewed By: patricklabatut

Differential Revision: D31895560

fbshipit-source-id: 1ecbb1782b0eafe132f88425c48487c2d0e10d2d
---
 tests/test_render_multigpu.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_render_multigpu.py b/tests/test_render_multigpu.py
index 40ff9f1d..a2bf7289 100644
--- a/tests/test_render_multigpu.py
+++ b/tests/test_render_multigpu.py
@@ -159,8 +159,8 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
         verts = ico_sphere(3).verts_padded()
         texs = verts.new_ones(verts.shape)
         model = Model()
+        model.to(GPU_LIST[0])
         model = nn.DataParallel(model, device_ids=GPU_LIST)
-        model.to(f"cuda:{model.device_ids[0]}")
 
         # Test a few iterations
         for _ in range(100):