def train_model(model, [criterion](https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss "torch.nn.CrossEntropyLoss"), optimizer, scheduler, num_epochs=25):
    since = time.time()

    # Create a temporary directory to save training checkpoints
    with TemporaryDirectory() as tempdir:
        best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')

        [torch.save](https://docs.pytorch.org/docs/stable/generated/torch.save.html#torch.save "torch.save")(model.state_dict(), best_model_params_path)
        best_acc = 0.0

        for epoch in range(num_epochs):
            print(f'Epoch {epoch}/{num_epochs - 1}')
            print('-' * 10)

            # Each epoch has a training and validation phase
            for phase in ['train', 'val']:
                if phase == 'train':
                    model.train()  # Set model to training mode
                else:
                    model.eval()   # Set model to evaluate mode

                running_loss = 0.0
                running_corrects = 0

                # Iterate over data.
                for [inputs](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"), labels in dataloaders[phase]:
                    [inputs](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor") = [inputs](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor").to(device)
                    labels = labels.to(device)

                    # zero the parameter gradients
                    optimizer.zero_grad()

                    # forward
                    # track history if only in train
                    with [torch.set_grad_enabled](https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.set_grad_enabled.html#torch.autograd.grad_mode.set_grad_enabled "torch.autograd.grad_mode.set_grad_enabled")(phase == 'train'):
                        outputs = model([inputs](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor"))
                        _, preds = [torch.max](https://docs.pytorch.org/docs/stable/generated/torch.max.html#torch.max "torch.max")(outputs, 1)
                        loss = [criterion](https://docs.pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html#torch.nn.CrossEntropyLoss "torch.nn.CrossEntropyLoss")(outputs, labels)

                        # backward + optimize only if in training phase
                        if phase == 'train':
                            loss.backward()
                            optimizer.step()

                    # statistics
                    running_loss += loss.item() * [inputs](https://docs.pytorch.org/docs/stable/tensors.html#torch.Tensor "torch.Tensor").size(0)
                    running_corrects += [torch.sum](https://docs.pytorch.org/docs/stable/generated/torch.sum.html#torch.sum "torch.sum")(preds == labels.data)
                if phase == 'train':
                    scheduler.step()

                epoch_loss = running_loss / dataset_sizes[phase]
                epoch_acc = running_corrects.double() / dataset_sizes[phase]

                print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

                # deep copy the model
                if phase == 'val' and epoch_acc > best_acc:
                    best_acc = epoch_acc
                    [torch.save](https://docs.pytorch.org/docs/stable/generated/torch.save.html#torch.save "torch.save")(model.state_dict(), best_model_params_path)

            print()

        time_elapsed = time.time() - since
        print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
        print(f'Best val Acc: {best_acc:4f}')

        # load best model weights
        model.load_state_dict([torch.load](https://docs.pytorch.org/docs/stable/generated/torch.load.html#torch.load "torch.load")(best_model_params_path, weights_only=True))
    return model