diff --git a/datadrivenpdes/pipelines/__init__.py b/datadrivenpdes/pipelines/__init__.py
index 5278647..547ed1f 100644
--- a/datadrivenpdes/pipelines/__init__.py
+++ b/datadrivenpdes/pipelines/__init__.py
@@ -1,3 +1,2 @@
 """Training and evaluation pipelines."""
-from datadrivenpdes.pipelines import beamlib
 from datadrivenpdes.pipelines import model_utils
diff --git a/datadrivenpdes/pipelines/beamlib.py b/datadrivenpdes/pipelines/beamlib.py
index 7634205..c9d2b07 100644
--- a/datadrivenpdes/pipelines/beamlib.py
+++ b/datadrivenpdes/pipelines/beamlib.py
@@ -14,6 +14,7 @@
 # limitations under the License.
 # ==============================================================================
 """Beam utilities."""
+
 import apache_beam as beam
 import numpy as np
 from typing import List, Tuple
diff --git a/datadrivenpdes/pipelines/create_training_data.py b/datadrivenpdes/pipelines/create_training_data.py
index a16d536..dffe7eb 100644
--- a/datadrivenpdes/pipelines/create_training_data.py
+++ b/datadrivenpdes/pipelines/create_training_data.py
@@ -19,12 +19,17 @@
 
 from absl import app
 from absl import flags
-import apache_beam as beam
+
+try:
+  import apache_beam as beam
+  from datadrivenpdes.pipelines import beamlib
+except:
+  pass
+
 import numpy as np
 from datadrivenpdes.core import builders
 from datadrivenpdes.core import equations
 from datadrivenpdes.core import grids
-from datadrivenpdes.pipelines import beamlib
 import tensorflow as tf
 
 # Ensure Equation subclasses are defined so we can look them up by name.
diff --git a/datadrivenpdes/pipelines/create_training_data_test.py b/datadrivenpdes/pipelines/create_training_data_test.py
index 97f2553..c1e90d9 100644
--- a/datadrivenpdes/pipelines/create_training_data_test.py
+++ b/datadrivenpdes/pipelines/create_training_data_test.py
@@ -19,7 +19,12 @@
 from absl import flags
 from absl.testing import flagsaver
 from absl.testing import parameterized
-import apache_beam as beam
+
+try:
+  import apache_beam as beam
+except:
+  pass
+
 from datadrivenpdes.core import builders
 from datadrivenpdes.pipelines import create_training_data
 from tensorflow import gfile
diff --git a/setup.py b/setup.py
index deb2835..8d45003 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,6 @@
 
 INSTALL_REQUIRES = [
     'absl-py',
-    'apache-beam',
     'numpy',
     'scipy',
     'xarray',