Create train.py

code/train.py

@@ -0,0 +1,28 @@
+import torch
+from tqdm import tqdm
+
+iterator = tqdm(dataloader, desc="Training", postfix={"train_loss":0.0})
+
+for item in iterator:
+  item = tokenizer.bos_token + " " + item[0] + " " + tokenizer.eos_token
+  encoded_inp = tokenizer(item, return_tensors='pt').input_ids.to("cuda")
+  logits = mamba_model(encoded_inp)
+
+  labels = encoded_inp.to(logits.device)
+  shift_logits = logits[:, :-1, :].contiguous()
+  labels = labels[:, 1:].contiguous()
+  loss_fct = torch.nn.CrossEntropyLoss()
+  loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), labels.view(-1))
+
+  optimizer.zero_grad(set_to_none=True)
+  loss.backward()
+  optimizer.step()
+
+  # moving data's from gpu to cpu
+  loss = loss.detach().cpu().numpy()
+  logits = logits.detach().cpu().numpy()
+  labels = labels.detach().cpu().numpy()
+  encoded_inp = encoded_inp.detach().cpu().numpy()
+  shift_logits = shift_logits.detach().cpu().numpy()
+
+  iterator.set_postfix({"train_loss": loss.item()}, refresh=False)