sparsify.blueprints.code_samples package
Submodules
sparsify.blueprints.code_samples.pytorch__integration module
sparsify.blueprints.code_samples.pytorch__training module
-
sparsify.blueprints.code_samples.pytorch__training.
train
(working_dir: str, config_path: str, model: torch.nn.modules.module.Module, train_dataset: torch.utils.data.dataset.Dataset, val_dataset: torch.utils.data.dataset.Dataset, batch_size: int, optim_const: Callable[torch.nn.modules.module.Module, torch.optim.optimizer.Optimizer], loss: Union[sparseml.pytorch.utils.loss.LossWrapper, Callable[[Any, Any], torch.Tensor]], devices: str)[source] Dataset setup