diff --git a/featurize_assignment.py b/featurize_assignment.py
index 3ccc1e6..4a2e2c4 100644
--- a/featurize_assignment.py
+++ b/featurize_assignment.py
@@ -29,8 +29,9 @@
 import numpy as np
 from ax.service.ax_client import AxClient, ObjectiveProperties
 
-from ax.modelbridge.factory import Models
-from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy
+from ax.modelbridge.factory import Generators
+from ax.generation_strategy.generation_strategy import GenerationStep, GenerationStrategy
+
 
 gs = GenerationStrategy(
     steps=[
@@ -66,4 +67,4 @@
 # TASK F: How many materials have a hardness greater than 43?
 # --------------------------------------------------------------------------------------
 
-# TODO: Your Code Goes Here
\ No newline at end of file
+# TODO: Your Code Goes Here
diff --git a/requirements.txt b/requirements.txt
index 28ac251..b2745ff 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,5 @@
 -e .  # Install local package in editable mode
-ax-platform>=0.3.0
+ax-platform<1
 matplotlib>=3.5.0
 CBFV>=0.1.0
 pytest>=7.0.0