OOM errors when running tests with bazel

Describe the bug
I had a lot of Out of Memory errors when running pruning tests with bazel.

This is what I tried on two different hardware (see below for configuration):

• bazel test tensorflow_model_optimization/python/core/sparsity/keras/...
  It gave me either a GPU->CPU Memcpy failed Fatal Python error: Aborted or a CUDA OOM on both hardware.
• bazel test tensorflow_model_optimization/python/core/sparsity/keras:prune_integration_test --runs_per_test=10
  Only one file to test for, but 10 runs. Same result.
• bazel test tensorflow_model_optimization/python/core/sparsity/keras:prune_integration_test
  Only one file to test for. It worked fine on hardware2, but did OOM on hardware1.
• bazel test tensorflow_model_optimization/python/core/sparsity/keras:prune_integration_test --runs_per_test=10 --jobs=1
  No parallel runs. It worked fine on hardware2, but did OOM on hardware1.

System information

Hardware configuration1: >32gb ram, >10cpus, Quadro P400 (2gb).

Hardware configuration2: >64gb ram, >30cpus, x4 GeForce GTX TITAN X.

TensorFlow version (installed from source or binary): tf2.4 and tf-nightly (tried both)

TensorFlow Model Optimization version (installed from source or binary): source.

Python version: 3.8

Software configuration: Ubuntu, Cuda11, Bazel3.7.2.

Describe the expected behavior
Tests should work.

Describe the current behavior
Tests are flaky due to memory issues.

Code to reproduce the issue
bazel test tensorflow_model_optimization/python/core/sparsity/keras/...

No source file or test implementation is named. Start by reproducing the issue with bazel test tensorflow_model_optimization/python/core/sparsity/keras:prune_integration_test, then compare the broader pruning test target and runs with --jobs=1 on the reported hardware configurations. Done means the pruning tests complete without GPU or CPU out-of-memory failures.