system.py

from typing import Any, Dict, Tuple, Union
import torch
from lighter import LighterModule

class LighterSystem(LighterModule):
    def training_step(self, batch: Any, batch_idx: int) -> Dict[str, torch.Tensor]:
        input, target = batch
        # Handle cases where SwaV needs epoch/step
        if hasattr(self.network, "forward") and "epoch" in self.network.forward.__code__.co_varnames:
            pred = self.network(input, epoch=self.current_epoch, step=self.global_step)
        else:
            pred = self.network(input)
        
        try:
            loss = self.criterion(pred, target)
        except TypeError:
            loss = self.criterion(pred)
            
        return {"loss": loss, "pred": pred, "target": target}

    def validation_step(self, batch: Any, batch_idx: int) -> Dict[str, torch.Tensor]:
        input, target = batch
        if hasattr(self.network, "forward") and "epoch" in self.network.forward.__code__.co_varnames:
            pred = self.network(input, epoch=self.current_epoch, step=self.global_step)
        else:
            pred = self.network(input)
        
        try:
            loss = self.criterion(pred, target)
        except TypeError:
            loss = self.criterion(pred)
            
        return {"loss": loss, "pred": pred, "target": target}

    def test_step(self, batch: Any, batch_idx: int) -> Dict[str, torch.Tensor]:
        input, target = batch
        pred = self.network(input)
        
        loss = None
        if self.criterion:
            try:
                loss = self.criterion(pred, target)
            except TypeError:
                loss = self.criterion(pred)
                
        return {"loss": loss, "pred": pred, "target": target}

    def predict_step(self, batch: Any, batch_idx: int) -> Any:
        if isinstance(batch, (list, tuple)) and len(batch) == 2:
            input, _ = batch
        else:
            input = batch
        return self.network(input)