code: def _train_with_dls(x, y):
      scaler = DynamicLossScaler()
      torch.manual_seed(42)
      lr = 1e-1
      n_epochs = 1000
      loss_fn = nn.MSELoss(reduction="mean")
      model = ManualLinearRegression().to(device)
      optimizer = optim.SGD(model.parameters(), lr=lr)
      for epoch in range(n_epochs):
          optimizer.zero_grad()
          model.train()
          yhat = model(x)
          loss = loss_fn(y, yhat)
          scaler.scale(loss).backward()
          scaler.step(optimizer)
          scaler.update()
      return model
  
file_path: fairscale/tests/experimental/optim/test_dynamic_loss_scaler.py
imports: import numpy
  import torch
  import torch.nn
  import torch.optim
  from fairscale.experimental.optim.dynamic_loss_scaler import DynamicLossScaler
label_desc: Use image node_id 2 for calling a global function with example usage: _train_with_dls(x, y) and returns: model