But when we initialize the optimizer as per above mentioned, optimizer could only take one step per input_tensor. Thus would you recommend taking multiple step for each input_tensor until a condition as shown below:
model.eval()
for inputs in train_dataloader:
     init_tensor = inputs["init_tensor"] # tensor to be optimized
     optimizer = optim.SGD(init_tensor, learning_rate)
     input_tensor = inputs["input_tensor"] # different images from loader
     input = init_tensor + input_tensor
     loss = 1
     while(loss > threshold):
          output = model(input)
          loss = loss_fn(output)
          model.zero_grad()
          loss.backward()
          optimizer.step()