From 9d0fc41ece04c51cf634c9fd15cb192af92576c9 Mon Sep 17 00:00:00 2001
From: Janice Lan <janlan@fb.com>
Date: Mon, 17 Jul 2023 14:28:13 -0700
Subject: [PATCH] remove old imports

---
 ocpmodels/tasks/task.py        | 1 -
 ocpmodels/trainers/__init__.py | 6 ++----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/ocpmodels/tasks/task.py b/ocpmodels/tasks/task.py
index 14d5cb4ac7..bf87a2f01d 100644
--- a/ocpmodels/tasks/task.py
+++ b/ocpmodels/tasks/task.py
@@ -9,7 +9,6 @@
 import os
 
 from ocpmodels.common.registry import registry
-from ocpmodels.trainers.forces_trainer import ForcesTrainer
 
 
 class BaseTask:
diff --git a/ocpmodels/trainers/__init__.py b/ocpmodels/trainers/__init__.py
index a93fc680bf..20b44d540c 100644
--- a/ocpmodels/trainers/__init__.py
+++ b/ocpmodels/trainers/__init__.py
@@ -5,10 +5,8 @@
 
 __all__ = [
     "BaseTrainer",
-    "ForcesTrainer",
-    "EnergyTrainer",
+    "OCPTrainer",
 ]
 
 from .base_trainer import BaseTrainer
-from .energy_trainer import EnergyTrainer
-from .forces_trainer import ForcesTrainer
+from .ocp_trainer import OCPTrainer