[test] Fix pooler unit tests on gpu (#1165)

Summary:
Pull Request resolved: #1165

Fix no attribute "cuda" runtime error when
testing poolers with cuda.
Unit tests pass on github actions as
those machines run without gpu.

Test Plan: Test unit tests with GPU on learnfair

Reviewed By: ebsmothers

Differential Revision: D32805045

Pulled By: Ryan-Qiyu-Jiang

fbshipit-source-id: 5d2272fc99688d71bbf9cd3b7b71c051deb1de36

Author: Ryan-Qiyu-Jiang

tests/modules/test_poolers.py

@@ -4,6 +4,7 @@
 
 import mmf.modules.poolers as poolers
 import torch
+from mmf.utils.general import get_current_device
 
 
 class TestModulePoolers(unittest.TestCase):
@@ -13,36 +14,29 @@ def setUp(self):
         self.num_tokens = 10
         self.embedding_size = 768
         self.token_len = 10
+        self.device = get_current_device()
         self.encoded_layers = [
-            torch.randn(self.batch_size, self.token_len, self.embedding_size),
-            torch.randn(self.batch_size, self.token_len, self.embedding_size),
-            torch.randn(self.batch_size, self.token_len, self.embedding_size),
+            torch.randn(self.batch_size, self.token_len, self.embedding_size).to(
+                self.device
+            )
+            for _ in range(3)
         ]
-        self.pad_mask = torch.randn(self.batch_size, self.token_len)
+        self.pad_mask = torch.randn(self.batch_size, self.token_len).to(self.device)
 
     def test_AverageConcat(self):
-        pool_fn = poolers.AverageConcatLastN(self.k)
+        pool_fn = poolers.AverageConcatLastN(self.k).to(self.device)
         out = pool_fn(self.encoded_layers, self.pad_mask)
-        if torch.cuda.is_available():
-            pool_fn.cuda()
-            out = pool_fn(self.encoded_layers.cuda(), self.pad_mask.cuda())
 
         assert torch.Size([self.batch_size, self.embedding_size * self.k]) == out.shape
 
     def test_AverageKFromLast(self):
-        pool_fn = poolers.AverageKFromLast(self.k)
+        pool_fn = poolers.AverageKFromLast(self.k).to(self.device)
         out = pool_fn(self.encoded_layers, self.pad_mask)
-        if torch.cuda.is_available():
-            pool_fn.cuda()
-            out = pool_fn([self.encoded_layers.cuda(), self.pad_mask.cuda()])
 
         assert torch.Size([self.batch_size, self.embedding_size]) == out.shape
 
     def test_AverageSumLastK(self):
-        pool_fn = poolers.AverageSumLastK(self.k)
+        pool_fn = poolers.AverageSumLastK(self.k).to(self.device)
         out = pool_fn(self.encoded_layers, self.pad_mask)
-        if torch.cuda.is_available():
-            pool_fn.cuda()
-            out = pool_fn([self.encoded_layers.cuda(), self.pad_mask.cuda()])
 
         assert torch.Size([self.batch_size, self.embedding_size]) == out.shape