Help with installation. No module named caffe2

Once I clone the repository and try to run ./test/dlrm_s_test.sh
I get some errors. This is both on my WSL setup and on a fresh new google Colab setup too.

Unable to import onnx or caffe2.python.onnx.frontend No module named 'caffe2'

Unable to import onnx. No module named 'onnx'

as well as

Traceback (most recent call last):
File "/home/jerry/workspace/embed/dlrm/dlrm_s_caffe2.py", line 87, in
from caffe2.proto import caffe2_pb2
ModuleNotFoundError: No module named 'caffe2'
Traceback (most recent call last):
File "/home/jerry/workspace/embed/dlrm/dlrm_s_caffe2.py", line 87, in
from caffe2.proto import caffe2_pb2
ModuleNotFoundError: No module named 'caffe2'
Traceback (most recent call last):
File "/home/jerry/workspace/embed/dlrm/dlrm_s_caffe2.py", line 87, in
from caffe2.proto import caffe2_pb2
ModuleNotFoundError: No module named 'caffe2'
Traceback (most recent call last):
File "/home/jerry/workspace/embed/dlrm/dlrm_s_caffe2.py", line 87, in
from caffe2.proto import caffe2_pb2
ModuleNotFoundError: No module named 'caffe2'

When I try to run pip install caffe2 to fix, it says I need tools.
Then when I install tools and try to run pip install caffe2 again, it says ImportError: No module named 'tools.setup_helpers'

I'm stuck. I want to use this repository to run a dlrm model and experiment with it.

It seems online that there are conflicting stuff about caffe2. What is it? Some places say it's already in pytorch, and that it's out of date or something.
I want to see what it shown in the README when I run ./test/dlrm_s_test.sh
and I also want to confirm I have the right output when I run ./bench/dlrm_s_benchmark.sh