code: def test_optim_with_scheduler(tmpdir, optim_fn, lr):
      max_epochs = 10
      layer = nn.Linear(10, 1)
      optimizer = optim_fn(layer.parameters(), lr=lr)
      scheduler = LinearWarmupCosineAnnealingLR(
          optimizer, warmup_epochs=2, max_epochs=max_epochs
      )
  
      for _ in range(max_epochs):
          dummy_input = torch.rand(1, 10)
          dummy_input.requires_grad = True
          result = layer(dummy_input)
          result.backward()
          optimizer.step()
          scheduler.step()
  
imports: import pytest
  import torch
  from flash.core.optimizers import LAMB, LARS, LinearWarmupCosineAnnealingLR
  from flash.core.utilities.imports import _TOPIC_CORE_AVAILABLE
  from torch import nn
label_desc: Use image node_id 2 for calling a global function with example usage: test_optim_with_scheduler(tmpdir, optim_fn, lr) without return types