flsim/utils/sample_model.py [74:89]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        )

    def fl_create_training_batch(self, **kwargs) -> None:
        return kwargs.get("batch", None)

    def fl_get_module(self) -> nn.Module:
        return self.model

    def fl_cuda(self) -> None:
        self.model = self.model.cuda()

    def get_eval_metrics(self, batch) -> FLBatchMetrics:
        with torch.no_grad():
            return self.fl_forward(batch)

    def get_num_examples(self, batch) -> int:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



flsim/utils/sample_model.py [272:287]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
        )

    def fl_create_training_batch(self, **kwargs) -> None:
        return kwargs.get("batch", None)

    def fl_get_module(self) -> nn.Module:
        return self.model

    def fl_cuda(self) -> None:
        self.model = self.model.cuda()

    def get_eval_metrics(self, batch) -> FLBatchMetrics:
        with torch.no_grad():
            return self.fl_forward(batch)

    def get_num_examples(self, batch) -> int:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -