Learning rate always zero with SequentialLR

jstrm: <h3>
<a name="problem-1" class="anchor" href="#problem-1"></a>Problem</h3>
<p>With the code below, I get a learning rate of zero for all iterations when using a small number of training samples, e.g., <code>batch_size=64</code>, <code>num_train_samples=74</code>, <code>num_epochs=10</code>, <code>warmup_epochs=2</code>. The milestone that I set seems to somehow be wrong.</p>
<h3>
<a name="what-i-tried-2" class="anchor" href="#what-i-tried-2"></a>What I tried</h3>
<p>The learning rate adapts as intended for larger trainings, e.g., <code>batch_size=64</code>, <code>num_train_samples=2328</code>, <code>num_epochs=2000</code>, <code>warmup_epochs=20</code> works fine.</p>
<p>Based on the print statements in <code>get_scheduler()</code> the variable values are as expected, so the variables seem to not be the problem here. Also, setting <code>warmup_epochs=0</code> results in the intended constant learning rate.</p>
<p>I suspect there is something wrong with how I set <code>warmup_steps</code> based on <code>batch_size</code>, <code>warmup_epochs</code> and <code>num_train_samples</code>, but I am at a loss what precisely is going wrong. Does anyone have an idea where my mistake is?</p>
<h3>
<a name="code-3" class="anchor" href="#code-3"></a>Code</h3>
<p>The code is part of a complex codebase, so I tried to provide only the relevant parts of the training process. I also included optimize_parameters, as I suspect there could be something wrong with the way I treat the <code>scaler</code>.</p>
<pre><code class="lang-auto">def warmup_wrapper(warmup_steps: int):
    """We need a closure here to set `warmup_steps`."""
    def warmup(current_step: int):
        # Linear warmup.
        return current_step / warmup_steps
    return warmup


def get_scheduler(optimizer, opt: argparse.Namespace):
    """Return a learning rate scheduler
    Parameters:
        optimizer          -- the optimizer of the network
        opt (option class) -- stores all the experiment flags
    """
        # We do scheduler.step() after each batch, so we need to calculate the actual number of warmup steps
        # based on the warmup_epochs given by the user.
    if opt.warmup_epochs &gt; 0:
        warmup_steps = opt.warmup_epochs * math.ceil(opt.num_train_samples / opt.batch_size)
        print(f"warmup_steps: {warmup_steps}")
        print(f"warmup_epochs: {opt.warmup_epochs}")
        print(f"train_samples: {opt.num_train_samples}")
    else:
        warmup_steps = 0
    main_scheduler = lr_scheduler.ConstantLR(optimizer, factor=1., total_iters=0, last_epoch=-1)
    if opt.warmup_epochs &gt; 0:        
        warmup_fn = warmup_wrapper(warmup_steps)
        warmup_scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_fn)
        scheduler = lr_scheduler.SequentialLR(
                                    optimizer, [warmup_scheduler, main_scheduler], milestones=[warmup_steps])
    else:
        scheduler = main_scheduler
    return scheduler

...

class VisionMethod():
    ...
    def optimize_parameters(
            self, scaler: torch.cuda.amp.GradScaler, use_mixed_precision: bool
     ) -&gt; Tuple[torch.cuda.amp.GradScaler, bool]:
        """Perform a forward pass, calculate the losses, and perform the backward pass for the current batch.
        """        
        device = 'cuda' if 'cuda' in str(self.device) else 'cpu'
        with torch.autocast(device_type=device, dtype=torch.float16, enabled=use_mixed_precision):            
            self.forward()
            self.calculate_losses() # sets self.total_loss
        self.set_requires_grad([self.segmentation_network], True)        
        if self.total_loss.requires_grad:
            scaler.scale(self.total_loss).backward()
        scaler.step(self.optimizer_segmentation)
        # We only want to update the learning rate schedulers, if the scaler was updated.
        # See https://github.com/pytorch/pytorch/issues/55585
        scaler_before = scaler.get_scale()
        scaler.update()
        scaler_after = scaler.get_scale()
        update_schedulers = False
        if scaler_before &lt;= scaler_after:
            update_schedulers = True
        self.optimizer_segmentation.zero_grad()
        return scaler, update_schedulers

def train_loop():
    ...
    method = VisionMethod()
    method.schedulers = get_scheduler(optimizer, self.opt)
    ...
    for epoch in range(next_epoch, opt.num_epochs + 1):
        method.train()
        iterations_in_epoch = 0

        scaler = torch.cuda.amp.GradScaler(enabled=opt.use_mixed_precision)

        with tqdm(
            total=len(train_dataset), desc='Training epoch {}/{}'.format(epoch, opt.num_epochs)
        ) as pbar:
            for i, data in enumerate(train_dataset):
                iterations_in_epoch += opt.batch_size
                total_iterations += opt.batch_size

                method.set_input(data)                 
                scaler, update_schedulers = method.optimize_parameters(scaler, opt.use_mixed_precision)
                # To prevent "UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`".
                if update_schedulers:
                    method.update_schedulers()
                del data

                if total_iterations % opt.print_freq &lt; opt.batch_size:        
                    learning_rate =method.optimizer.param_groups[0]['lr'])                    
                    method.tensorboardnetworkwriter.add_scalar("Learning rate_{}".format(method.network_name), learning_rate, iteration)

                pbar.update(opt.batch_size)
</code></pre>