Repository: mosaicml/composer

Issue:
Segmentation Fault in CPU Test
** Environment **
<!-- Provide the below as necessary -->

- OS: [Ubuntu 20.04]
- Hardware (GPU, or instance type): [GEFORCE 3080]

<!-- if your bug is CUDA related, please also provide your CUDA version -->
Not CUDA related

** To reproduce

Steps to reproduce the behavior:

1. make test

## Expected behavior

<!-- A clear and concise description of what you would expect to happen. -->
Tests should pass

## Additional context

<!-- Please provide an additional context. -->
Hi there! I am preparing my application for the research engineer role at Mosaic, and I thought I would start by contributing to this awesome repo! After I install dependencies and run the CPU test, I receive the following error. This error is not CUDA related because GPU and distributed GPU tests both pass, but not CPU or distributed CPU tests. The error seems to be pointing to the line `trainer.fit()`, but I have also made sure that the trainer runs properly when I run normal training. That's why I am confused as to why this error would happen.

```
▶ make test                                    
python   -m pytest    --duration all  
=================================================================== test session starts ====================================================================
platform linux -- Python 3.8.12, pytest-7.1.0, pluggy-1.0.0
rootdir: /workspace/composer, configfile: pyproject.toml
plugins: timeout-2.1.0
timeout: 2.0s
timeout method: signal
timeout func_only: False
collected 1659 items / 223 deselected / 1436 selected                                                                                                      

tests/test_dataset_registry.py .x...xxxxsssssss                                                                                                      [  1%]
tests/test_device.py ......                                                                                                                          [  1%]
tests/test_engine.py ...................................................................................................                             [  8%]
tests/test_events.py .......................                                                                                                         [ 10%]
tests/test_hparams.py ..........s.................s...................xx.xxxx.sx....xxx.x.xx..xxs........xxxxxxxxx                                   [ 16%]
tests/test_iter_helpers.py ..........                                                                                                                [ 17%]
tests/test_load.py ..ss.ssss.ss....sss.s.ss..sss...............                                                                                      [ 20%]
tests/test_loss.py ............xxxxxxXXXXXX............                                                                                              [ 22%]
tests/test_model_registry.py ............                                                                                                            [ 23%]
tests/test_object_store.py .                                                                                                                         [ 23%]
tests/test_pillow_simd.py s                                                                                                                          [ 23%]
tests/test_state.py .                                                                                                                                [ 23%]
tests/test_string_enum.py ....                                                                                                                       [ 24%]
tests/test_surgery.py ..........                                                                                                                     [ 24%]
tests/test_time.py .............                                                                                                                     [ 25%]
tests/algorithms/test_algorithm_registry.py ......................                                                                                   [ 27%]
tests/algorithms/test_augmentations_functional.py ....................................................................................               [ 33%]
tests/algorithms/test_blurpool.py ..................................................................                                                 [ 37%]
tests/algorithms/test_blurpool_layers.py ............................................................................................                [ 44%]
tests/algorithms/test_channels_last.py ..........                                                                                                    [ 44%]
tests/algorithms/test_colout.py ...................................................                                                                  [ 48%]
tests/algorithms/test_cutmix.py ..........                                                                                                           [ 48%]
tests/algorithms/test_cutout.py ..............................................................................................                       [ 55%]
tests/algorithms/test_factorize_algorithm.py ............                                                                                            [ 56%]
tests/algorithms/test_factorize_core.py ................................                                                                             [ 58%]
tests/algorithms/test_factorized_modules.py ................................................................................                         [ 64%]
tests/algorithms/test_ghost_batchnorm.py ....xx.................................................                                                     [ 67%]
tests/algorithms/test_label_smoothing.py .............................................................                                               [ 72%]
tests/algorithms/test_layer_freezing.py ...                                                                                                          [ 72%]
tests/algorithms/test_mixup.py ....                                                                                                                  [ 72%]
tests/algorithms/test_progressive_resizing.py .................................................                                                      [ 76%]
tests/algorithms/test_selective_backprop.py .......x......................                                                                           [ 78%]
tests/algorithms/test_squeeze_excite.py .....                                                                                                        [ 78%]
tests/algorithms/test_stochastic_depth.py .................................                                                                          [ 80%]
tests/callbacks/test_callbacks.py ........................                                                                                           [ 82%]
tests/callbacks/test_grad_monitor.py ..                                                                                                              [ 82%]
tests/callbacks/test_memory_monitor.py .                                                                                                             [ 82%]
tests/callbacks/test_speed_monitor.py .                                                                                                              [ 82%]
tests/datasets/test_add_dataset_transform.py .......                                                                                                 [ 83%]
tests/datasets/test_c4.py ssss                                                                                                                       [ 83%]
tests/datasets/test_dataset_utils.py ..                                                                                                              [ 83%]
tests/datasets/test_segmentation_transforms.py .......                                                                                               [ 84%]
tests/datasets/test_synthetic_data.py .s.......s..                                                                                                   [ 85%]
tests/datasets/test_synthetic_lm_data.py ......                                                                                                      [ 85%]
tests/loggers/test_file_logger.py ...                                                                                                                [ 85%]
tests/loggers/test_in_memory_logger.py ..                                                                                                            [ 85%]
tests/loggers/test_logger.py .                                                                                                                       [ 85%]
tests/loggers/test_object_store_logger.py ...                                                                                                        [ 86%]
tests/loggers/test_progress_bar_logger.py .                                                                                                          [ 86%]
tests/loggers/test_wandb_logger.py .                                                                                                                 [ 86%]
tests/metrics/test_metrics.py .                                                                                                                      [ 86%]
tests/metrics/test_nlp_metrics.py ......................                                                                                             [ 87%]
tests/profiler/test_json_trace_handler.py Fatal Python error: Segmentation fault

Thread 0x00007fb39d949700 (most recent call first):
  File "/opt/conda/envs/jax/lib/python3.8/threading.py", line 306 in wait
  File "/opt/conda/envs/jax/lib/python3.8/threading.py", line 558 in wait
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/tqdm/_monitor.py", line 60 in run
  File "/opt/conda/envs/jax/lib/python3.8/threading.py", line 932 in _bootstrap_inner
  File "/opt/conda/envs/jax/lib/python3.8/threading.py", line 890 in _bootstrap

Thread 0x00007fb3a77ed700 (most recent call first):
<no Python frame>

Current thread 0x00007fb4f62744c0 (most recent call first):
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/torch/autograd/profiler.py", line 238 in export_chrome_trace
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/torch/profiler/profiler.py", line 121 in export_chrome_trace
  File "/workspace/composer/composer/profiler/torch_profiler.py", line 209 in handler_fn
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/torch/profiler/profiler.py", line 487 in _trace_ready
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/torch/profiler/profiler.py", line 493 in _transit_action
  File "/opt/conda/envs/jax/lib/python3.8/site-packages/torch/profiler/profiler.py", line 479 in step
  File "/workspace/composer/composer/profiler/torch_profiler.py", line 245 in batch_end
  File "/workspace/composer/composer/core/callback.py", line 99 in run_event
  File "/workspace/composer/composer/core/engine.py", line 292 in _run_callbacks
  File "/workspace/composer/composer/core/engine.py", line 185 in run_event
  File "/workspace/composer/composer/trainer/trainer.py", line 1087 in _train_loop
  File "/workspace/composer/composer/trainer/trainer.py", line 881 in fit
  File "/workspace/composer/tests/profiler/test_json_trace_handler.py", line 32 in test_json_trace_profiler_handler
  File "/root/.local/lib/python3.8/site-packages/_pytest/python.py", line 192 in pytest_pyfunc_call
  File "/root/.local/lib/python3.8/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/root/.local/lib/python3.8/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/root/.local/lib/python3.8/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/root/.local/lib/python3.8/site-packages/_pytest/python.py", line 1761 in runtest
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 166 in pytest_runtest_call
  File "/root/.local/lib/python3.8/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/root/.local/lib/python3.8/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/root/.local/lib/python3.8/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 259 in <lambda>
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 338 in from_call
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 258 in call_runtest_hook
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 219 in call_and_report
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 130 in runtestprotocol
  File "/root/.local/lib/python3.8/site-packages/_pytest/runner.py", line 111 in pytest_runtest_protocol
  File "/root/.local/lib/python3.8/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/root/.local/lib/python3.8/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/root/.local/lib/python3.8/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/root/.local/lib/python3.8/site-packages/_pytest/main.py", line 347 in pytest_runtestloop
  File "/root/.local/lib/python3.8/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/root/.local/lib/python3.8/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/root/.local/lib/python3.8/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/root/.local/lib/python3.8/site-packages/_pytest/main.py", line 322 in _main
  File "/root/.local/lib/python3.8/site-packages/_pytest/main.py", line 268 in wrap_session
  File "/root/.local/lib/python3.8/site-packages/_pytest/main.py", line 315 in pytest_cmdline_main
  File "/root/.local/lib/python3.8/site-packages/pluggy/_callers.py", line 39 in _multicall
  File "/root/.local/lib/python3.8/site-packages/pluggy/_manager.py", line 80 in _hookexec
  File "/root/.local/lib/python3.8/site-packages/pluggy/_hooks.py", line 265 in __call__
  File "/root/.local/lib/python3.8/site-packages/_pytest/config/__init__.py", line 164 in main
  File "/root/.local/lib/python3.8/site-packages/_pytest/config/__init__.py", line 187 in console_main
  File "/root/.local/lib/python3.8/site-packages/pytest/__main__.py", line 5 in <module>
  File "/opt/conda/envs/jax/lib/python3.8/runpy.py", line 87 in _run_code
  File "/opt/conda/envs/jax/lib/python3.8/runpy.py", line 194 in _run_module_as_main
make: *** [Makefile:30: test] Segmentation fault (core dumped)
```


Discussion:
Hi there! Thanks for catching this issue. 

Which PyTorch version are you using, and whether you are using a custom build of Python? It seems like the failing test is related to the profiler.

In addition, if you comment out the profiler tests, does the rest of the test suite pass?
I think it's something to do with the file test_json_trace_handler.py, because when I comment everything inside that file this is the result I get. I emphasize that I didn't even touch test_profiler.py.
```
============================= 1319 passed, 61 skipped, 223 deselected, 49 xfailed, 6 xpassed, 327 warnings in 78.53s (0:01:18) =============================
```

I use a Conda environment, and this is my versions
```
_libgcc_mutex             0.1                 conda_forge    conda-forge
_openmp_mutex             4.5                       1_gnu    conda-forge
absl-py                   1.0.0                    pypi_0    pypi
aiohttp                   3.8.1                    pypi_0    pypi
aiosignal                 1.2.0                    pypi_0    pypi
aom                       3.3.0                h27087fc_1    conda-forge
asdfghjkl                 0.1a2                    pypi_0    pypi
async-timeout             4.0.2                    pypi_0    pypi
backpack-for-pytorch      1.5.0                    pypi_0    pypi
blas                      1.0                         mkl    defaults
brotlipy                  0.7.0           py38h27cfd23_1003    defaults
bzip2                     1.0.8                h7b6447c_0    defaults
ca-certificates           2022.3.18            h06a4308_0    defaults
cachetools                5.0.0                    pypi_0    pypi
certifi                   2021.10.8        py38h06a4308_2    defaults
cffi                      1.15.0           py38hd667e15_1    defaults
charset-normalizer        2.0.12             pyhd8ed1ab_0    conda-forge
click                     8.1.0                    pypi_0    pypi
cryptography              36.0.2           py38h2b5fc30_0    conda-forge
cudatoolkit               11.3.1               h2bc3f7f_2    defaults
cycler                    0.11.0                   pypi_0    pypi
datasets                  1.18.4                   pypi_0    pypi
deepspeed                 0.5.10                   pypi_0    pypi
dill                      0.3.4                    pypi_0    pypi
docker-pycreds            0.4.0                    pypi_0    pypi
einops                    0.4.1                    pypi_0    pypi
ffmpeg                    5.0.0                h594f047_1    conda-forge
filelock                  3.6.0                    pypi_0    pypi
fonttools                 4.31.2                   pypi_0    pypi
freetype                  2.11.0               h70c0345_0    defaults
frozenlist                1.3.0                    pypi_0    pypi
fsspec                    2022.2.0                 pypi_0    pypi
giflib                    5.2.1                h7b6447c_0    defaults
gitdb                     4.0.9                    pypi_0    pypi
gitpython                 3.1.27                   pypi_0    pypi
gmp                       6.2.1                h2531618_2    defaults
gnutls                    3.6.15               he1e5248_0    defaults
google-auth               2.6.2                    pypi_0    pypi
google-auth-oauthlib      0.4.6                    pypi_0    pypi
grpcio                    1.45.0                   pypi_0    pypi
hjson                     3.0.2                    pypi_0    pypi
huggingface-hub           0.4.0                    pypi_0    pypi
icu                       69.1                 h9c3ff4c_0    conda-forge
idna                      3.3                pyhd3eb1b0_0    defaults
importlib-metadata        4.11.3                   pypi_0    pypi
intel-openmp              2021.4.0          h06a4308_3561    defaults
jbig                      2.1                  hdba287a_0    defaults
joblib                    1.1.0                    pypi_0    pypi
jpeg                      9e                   h7f98852_0    conda-forge
kiwisolver                1.4.0                    pypi_0    pypi
lame                      3.100                h7b6447c_0    defaults
laplace-torch             0.1a2                    pypi_0    pypi
lcms2                     2.12                 h3be6417_0    defaults
ld_impl_linux-64          2.36.1               hea4e1c9_2    conda-forge
lerc                      3.0                  h295c915_0    defaults
libdeflate                1.10                 h7f98852_0    conda-forge
libdrm                    2.4.109              h7f98852_0    conda-forge
libffi                    3.3                  he6710b0_2    defaults
libgcc-ng                 11.2.0              h1d223b6_14    conda-forge
libgomp                   11.2.0              h1d223b6_14    conda-forge
libiconv                  1.16                 h516909a_0    conda-forge
libidn2                   2.3.2                h7f8727e_0    defaults
libpciaccess              0.16                 h516909a_0    conda-forge
libpng                    1.6.37               hbc83047_0    defaults
libstdcxx-ng              11.2.0              he4da1e4_14    conda-forge
libtasn1                  4.18.0               h7f98852_0    conda-forge
libtiff                   4.3.0                h542a066_3    conda-forge
libunistring              0.9.10               h27cfd23_0    defaults
libuv                     1.43.0               h7f98852_0    conda-forge
libva                     2.14.0               h7f98852_0    conda-forge
libvpx                    1.11.0               h295c915_0    defaults
libwebp                   1.2.2                h55f646e_0    defaults
libwebp-base              1.2.2                h7f8727e_0    defaults
libxcb                    1.14                 h7b6447c_0    defaults
libxml2                   2.9.12               h885dcf4_1    conda-forge
libzlib                   1.2.11            h36c2ea0_1013    conda-forge
lz4-c                     1.9.3                h295c915_1    defaults
markdown                  3.3.6                    pypi_0    pypi
matplotlib                3.5.1                    pypi_0    pypi
mkl                       2021.4.0           h06a4308_640    defaults
mkl-service               2.4.0            py38h7f8727e_0    defaults
mkl_fft                   1.3.1            py38hd3c417c_0    defaults
mkl_random                1.2.2            py38h51133e4_0    defaults
monai                     0.8.1                    pypi_0    pypi
multidict                 6.0.2                    pypi_0    pypi
multiprocess              0.70.12.2                pypi_0    pypi
ncurses                   6.3                  h7f8727e_2    defaults
netcal                    1.2.1                    pypi_0    pypi
nettle                    3.7.3                hbbd107a_1    defaults
ninja                     1.10.2.3                 pypi_0    pypi
numpy                     1.21.2           py38h20f2e39_0    defaults
numpy-base                1.21.2           py38h79a1101_0    defaults
oauthlib                  3.2.0                    pypi_0    pypi
openh264                  2.1.1                h4ff587b_0    defaults
openssl                   1.1.1n               h7f8727e_0    defaults
opt-einsum                3.3.0                    pypi_0    pypi
packaging                 21.3                     pypi_0    pypi
pandas                    1.4.1                    pypi_0    pypi
pathtools                 0.1.2                    pypi_0    pypi
pillow                    9.0.1            py38h22f2fdc_0    defaults
pip                       22.0.4             pyhd8ed1ab_0    conda-forge
promise                   2.3                      pypi_0    pypi
protobuf                  3.19.4                   pypi_0    pypi
py-cpuinfo                8.0.0                    pypi_0    pypi
pyarrow                   7.0.0                    pypi_0    pypi
pyasn1                    0.4.8                    pypi_0    pypi
pyasn1-modules            0.2.8                    pypi_0    pypi
pycocotools               2.0.4                    pypi_0    pypi
pycparser                 2.21               pyhd3eb1b0_0    defaults
pyopenssl                 22.0.0             pyhd3eb1b0_0    defaults
pyparsing                 3.0.7                    pypi_0    pypi
pyro-api                  0.1.2                    pypi_0    pypi
pyro-ppl                  1.8.1                    pypi_0    pypi
pysocks                   1.7.1            py38h06a4308_0    defaults
python                    3.8.12               h12debd9_0    defaults
python-dateutil           2.8.2                    pypi_0    pypi
python_abi                3.8                      2_cp38    conda-forge
pytorch                   1.11.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
pytorch-mutex             1.0                        cuda    pytorch
pyyaml                    6.0                      pypi_0    pypi
readline                  8.1.2                h7f8727e_1    defaults
regex                     2022.3.15                pypi_0    pypi
requests                  2.27.1             pyhd3eb1b0_0    defaults
requests-oauthlib         1.3.1                    pypi_0    pypi
responses                 0.18.0                   pypi_0    pypi
rsa                       4.8                      pypi_0    pypi
sacremoses                0.0.49                   pypi_0    pypi
scikit-learn              1.0.2                    pypi_0    pypi
scipy                     1.8.0                    pypi_0    pypi
sentry-sdk                1.5.8                    pypi_0    pypi
setproctitle              1.2.2                    pypi_0    pypi
setuptools                60.10.0          py38h578d9bd_0    conda-forge
shortuuid                 1.0.8                    pypi_0    pypi
six                       1.16.0             pyhd3eb1b0_1    defaults
smmap                     5.0.0                    pypi_0    pypi
sqlite                    3.38.0               hc218d9a_0    defaults
svt-av1                   0.9.1                h27087fc_0    conda-forge
tensorboard               2.8.0                    pypi_0    pypi
tensorboard-data-server   0.6.1                    pypi_0    pypi
tensorboard-plugin-wit    1.8.1                    pypi_0    pypi
termcolor                 1.1.0                    pypi_0    pypi
threadpoolctl             3.1.0                    pypi_0    pypi
tikzplotlib               0.10.1                   pypi_0    pypi
timm                      0.5.4                    pypi_0    pypi
tk                        8.6.12               h27826a3_0    conda-forge
tokenizers                0.11.6                   pypi_0    pypi
torchaudio                0.11.0               py38_cu113    pytorch
torchvision               0.12.0               py38_cu113    pytorch
tqdm                      4.63.1                   pypi_0    pypi
transformers              4.17.0                   pypi_0    pypi
triton                    1.0.0                    pypi_0    pypi
typing_extensions         4.1.1              pyh06a4308_0    defaults
urllib3                   1.26.9             pyhd8ed1ab_0    conda-forge
vit-pytorch               0.27.0                   pypi_0    pypi
wandb                     0.12.11                  pypi_0    pypi
webcolors                 1.11.1                   pypi_0    pypi
werkzeug                  2.0.3                    pypi_0    pypi
wheel                     0.37.1             pyhd3eb1b0_0    defaults
wurlitzer                 3.0.2                    pypi_0    pypi
x264                      1!161.3030           h7f98852_1    conda-forge
x265                      3.5                  h924138e_3    conda-forge
xorg-fixesproto           5.0               h7f98852_1002    conda-forge
xorg-kbproto              1.0.7             h7f98852_1002    conda-forge
xorg-libx11               1.7.2                h7f98852_0    conda-forge
xorg-libxext              1.3.4                h7f98852_1    conda-forge
xorg-libxfixes            5.0.3             h7f98852_1004    conda-forge
xorg-xextproto            7.3.0             h7f98852_1002    conda-forge
xorg-xproto               7.0.31            h27cfd23_1007    defaults
xxhash                    3.0.0                    pypi_0    pypi
xz                        5.2.5                h7b6447c_0    defaults
yacs                      0.1.8                    pypi_0    pypi
yarl                      1.7.2                    pypi_0    pypi
yaspin                    2.1.0                    pypi_0    pypi
zipp                      3.7.0                    pypi_0    pypi
zlib                      1.2.11            h36c2ea0_1013    conda-forge
zstd                      1.5.2                ha95c52a_0    conda-forge
```
Thanks for this info. The JSON trace handler test uses the PyTorch profiler underneath. We have not yet tested on PyTorch 1.11, so a change to the PyTorch profiler may have broken this test. We are looking into it.
Ah good point. I believe it is because of torch1.11, because this problem disappears with torch1.9. I will leave the issue open for now

Resolving patch:
diff --git a/composer/trainer/trainer_hparams.py b/composer/trainer/trainer_hparams.py
--- a/composer/trainer/trainer_hparams.py
+++ b/composer/trainer/trainer_hparams.py
@@ -409,7 +409,7 @@ class TrainerHparams(hp.Hparams):
     )
     torch_prof_profile_memory: bool = hp.optional(
         "Track tensor memory allocations and frees. Ignored if `prof_trace_handlers` is not specified.",
-        default=True,
+        default=False,
     )
     torch_prof_with_stack: bool = hp.optional(
         "Record stack information. Ignored if `prof_trace_handlers` is not specified.",
@@ -417,7 +417,7 @@ class TrainerHparams(hp.Hparams):
     )
     torch_prof_with_flops: bool = hp.optional(
         "Estimate flops for operators. Ignored if `prof_trace_handlers` is not specified.",
-        default=True,
+        default=False,
     )
 
     def validate(self):


Test patch:
diff --git a/tests/profiler/test_json_trace_handler.py b/tests/profiler/test_json_trace_handler.py
--- a/tests/profiler/test_json_trace_handler.py
+++ b/tests/profiler/test_json_trace_handler.py
@@ -10,6 +10,7 @@
 from composer.trainer import TrainerHparams
 
 
+# This test shouldn't run with the Torch profiler enabled, not providing a model or data can cause a seg fault
 @pytest.mark.timeout(10)
 def test_json_trace_profiler_handler(composer_trainer_hparams: TrainerHparams, tmpdir: pathlib.Path):
     profiler_file = os.path.join(tmpdir, 'trace.json')
@@ -27,6 +28,10 @@ def test_json_trace_profiler_handler(composer_trainer_hparams: TrainerHparams, t
     composer_trainer_hparams.sys_prof_net = False
     composer_trainer_hparams.sys_prof_disk = False
     composer_trainer_hparams.sys_prof_memory = False
+    composer_trainer_hparams.torch_prof_record_shapes = False
+    composer_trainer_hparams.torch_prof_profile_memory = False
+    composer_trainer_hparams.torch_prof_with_stack = False
+    composer_trainer_hparams.torch_prof_with_flops = False
 
     trainer = composer_trainer_hparams.initialize_object()
     trainer.fit()
diff --git a/tests/test_precision.py b/tests/test_precision.py
--- a/tests/test_precision.py
+++ b/tests/test_precision.py
@@ -39,6 +39,7 @@ def run_and_measure_memory(precision: Precision) -> int:
 
 
 @pytest.mark.gpu
+@pytest.mark.timeout(5)
 @pytest.mark.parametrize("precision", [Precision.AMP, Precision.BF16])
 def test_precision_memory(precision):
     if version.parse(torch.__version__) < version.parse("1.10"):