Loss not going down - first-time user

[alexlewisroberts]

I validated that the loss decreases without using Lightning. Here is my code:

train_ds = TensorDataset(X_train,y_train)
batch_size = 10000
train_dl = DataLoader(train_ds, batch_size, drop_last = False)
class Linear(pl.LightningModule):
  def __init__(self,input_size,output_size):
  super().__init__()
  self.model = nn.Linear(input_size,output_size)
  self.loss = nn.MSELoss(reduction=‘mean’)

def forward(self,x):
  return self.model(x)

def training_step(self,batch,batch_idx):
  x,y = batch
  loss = self.loss(self(x), y) 
  return loss

def configure_optimizers(self):
  return torch.optim.SGD(self.parameters(),lr=0.1)

input_size = 21601
output_size = 140
model = Linear(input_size, output_size)
trainer = pl.Trainer(max_epochs = 1, log_every_n_steps=1)
trainer.fit(model=model,train_dataloaders=train_dl)

[akihironitta]

Have you tried printing the loss in your training_step or increasing max_epochs?

def training_step(self,batch,batch_idx):
  x,y = batch
  loss = self.loss(self(x), y) 
+ print(loss)
  return loss

Also, you might want to check our examples here: https://github.com/Lightning-AI/lightning/tree/c39c8eb2e4f49e35ef38a2ae6de2765ec71623db/examples