def step(self, batch: any, batch_idx: BatchIndex):
    '\n        ### Training/validation step\n        '
    (data, target) = (batch[0].to(self.device), batch[1].to(self.device))
    if self.mode.is_train:
        tracker.add_global_step((data.shape[0] * data.shape[1]))
    with self.mode.update(is_log_activations=batch_idx.is_last):
        mem = self.memory.get()
        (output, new_mem) = self.model(data, mem)
        (mem, mem_to_compress) = self.merge_compress_memory(mem, new_mem)
        self.memory.set(mem)
    loss = self.loss_func(output, target)
    tracker.add('loss.', loss)
    if mem_to_compress:
        ar_loss = self.attention_reconstruction_loss(new_mem, mem_to_compress)
        tracker.add('ar_loss.', ar_loss)
        loss = (loss + ar_loss)
    self.accuracy(output, target)
    self.accuracy.track()
    if self.mode.is_train:
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
        self.optimizer.step()
        if batch_idx.is_last:
            tracker.add('model', self.model)
        self.optimizer.zero_grad()
    tracker.save()
### Training/validation step
def step(self, batch: any, batch_idx: BatchIndex):
    '\n        \n        '
    (data, target) = (batch[0].to(self.device), batch[1].to(self.device))
    if self.mode.is_train:
        tracker.add_global_step((data.shape[0] * data.shape[1]))
    with self.mode.update(is_log_activations=batch_idx.is_last):
        mem = self.memory.get()
        (output, new_mem) = self.model(data, mem)
        (mem, mem_to_compress) = self.merge_compress_memory(mem, new_mem)
        self.memory.set(mem)
    loss = self.loss_func(output, target)
    tracker.add('loss.', loss)
    if mem_to_compress:
        ar_loss = self.attention_reconstruction_loss(new_mem, mem_to_compress)
        tracker.add('ar_loss.', ar_loss)
        loss = (loss + ar_loss)
    self.accuracy(output, target)
    self.accuracy.track()
    if self.mode.is_train:
        loss.backward()
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.grad_norm_clip)
        self.optimizer.step()
        if batch_idx.is_last:
            tracker.add('model', self.model)
        self.optimizer.zero_grad()
    tracker.save()