Lint

Author: Jack-Khuu

src/forge/actors/reference_model.py

@@ -5,28 +5,20 @@
 # LICENSE file in the root directory of this source tree.
 
 
-import asyncio
 import logging
 import math
 import os
 
-from collections import deque
 from collections.abc import Mapping
 from dataclasses import dataclass, field, fields
 
-from typing import Any
-
 import torch
 from monarch.actor import current_rank, current_size, endpoint
-from omegaconf import DictConfig, OmegaConf
-from torch import nn
 
-from torchtitan.components.lr_scheduler import LRSchedulersContainer
-from torchtitan.config.job_config import Compile, Checkpoint, Comm, Model, Parallelism
-from torchtitan.distributed import ParallelDims, utils as dist_utils
+from torchtitan.config.job_config import Compile, Checkpoint, Model, Parallelism
+from torchtitan.distributed import utils as dist_utils
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
-from transformers import AutoModelForCausalLM
 
 from forge.controller import ForgeActor
 

tests/unit_tests/actors/test_reference_actor.py

@@ -11,7 +11,7 @@
 import pytest
 import torch
 
-from forge.actors.reference_actor import compute_logprobs
+from forge.actors.reference_model import compute_logprobs
 
 
 class TestComputeLogprobs: