diff --git a/reagent/mab/thompson_sampling.py b/reagent/mab/thompson_sampling.py
index a61f1330f..72967d59b 100644
--- a/reagent/mab/thompson_sampling.py
+++ b/reagent/mab/thompson_sampling.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
-
+# noqa
 from abc import abstractmethod
 from typing import Optional, List
 
diff --git a/reagent/test/preprocessing/test_transforms.py b/reagent/test/preprocessing/test_transforms.py
index b63391487..57e9542d2 100644
--- a/reagent/test/preprocessing/test_transforms.py
+++ b/reagent/test/preprocessing/test_transforms.py
@@ -9,6 +9,7 @@
 import numpy as np
 import reagent.core.types as rlt
 import torch
+import fbgemm_gpu
 from reagent.preprocessing import transforms
 from reagent.preprocessing.types import InputColumn
 
diff --git a/setup.cfg b/setup.cfg
index 7d3612167..019b8804d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -17,6 +17,7 @@ install_requires =
   # issue: https://github.com/openai/spinningup/issues/178
   cloudpickle~=1.2.0
   iopath
+  fbgemm-gpu-nightly
   numpy>=1.17.2
   pandas>=1.0.3
   # https://github.com/samuelcolvin/pydantic/issues/2042
@@ -52,7 +53,9 @@ ax = ax-platform
 
 lite = nevergrad>=0.4.3
 
-torchrec_gpu = torchrec-nightly
+torchrec_gpu =
+  fbgemm-gpu-nightly
+  torchrec-nightly
 
 torchrec_cpu = torchrec-nightly-cpu