docstr: For float16 only.
  Sub-sample sentences in a batch, and add padding, so that each dimension is a multiple of 8.
  
  Input:
  ------
      x: `torch.tensor(bs, seq_length)` - The token ids.
      lengths: `torch.tensor(bs, seq_length)` - The lengths of each of the sequence in the batch.
  
  Output:
  -------
      x:  `torch.tensor(new_bs, new_seq_length)` - The updated token ids.
      lengths: `torch.tensor(new_bs, new_seq_length)` - The updated lengths.
docstr_tok: ["For","float16","only",".","Sub-sample","sentences","in","a","batch",",","and","add","padding",",","so","that","each","dimension","is","a","multiple","of","8",".","Input",":","--","--","--","x",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","token","ids",".","lengths",":","`","torch.tensor","(","bs",",","seq_length",")","`","-","The","lengths","of","each","of","the","sequence","in","the","batch",".","Output",":","--","--","--","-","x",":","`","torch.tensor","(","new_bs",",","new_seq_length",")","`","-","The","updated","token","ids",".","lengths",":","`","torch.tensor","(","new_bs",",","new_seq_length",")","`","-","The","updated","lengths","."]
code: def round_batch(self, x: torch.tensor, lengths: torch.tensor):
      """
      For float16 only.
      Sub-sample sentences in a batch, and add padding, so that each dimension is a multiple of 8.
  
      Input:
      ------
          x: `torch.tensor(bs, seq_length)` - The token ids.
          lengths: `torch.tensor(bs, seq_length)` - The lengths of each of the sequence in the batch.
  
      Output:
      -------
          x:  `torch.tensor(new_bs, new_seq_length)` - The updated token ids.
          lengths: `torch.tensor(new_bs, new_seq_length)` - The updated lengths.
      """
      if not self.fp16 or len(lengths) < 8:
          return x, lengths
  
      # number of sentences == 0 [8]
      bs1 = len(lengths)
      bs2 = 8 * (bs1 // 8)
      assert bs2 > 0 and bs2 % 8 == 0
      if bs1 != bs2:
          idx = torch.randperm(bs1)[:bs2]
          lengths = lengths[idx]
          slen = lengths.max().item()
          x = x[idx, :slen]
      else:
          idx = None
  
      # sequence length == 0 [8]
      ml1 = x.size(1)
      if ml1 % 8 != 0:
          pad = 8 - (ml1 % 8)
          ml2 = ml1 + pad
          if self.mlm:
              pad_id = self.params.special_tok_ids["pad_token"]
          else:
              pad_id = self.params.special_tok_ids["unk_token"]
          padding_tensor = torch.zeros(
              bs2, pad, dtype=torch.long, device=x.device
          ).fill_(pad_id)
          x = torch.cat([x, padding_tensor], 1)
          assert x.size() == (bs2, ml2)
  
      assert x.size(0) % 8 == 0
      assert x.size(1) % 8 == 0
      return x, lengths
  
imports: import math
  import os
  import time
  import psutil
  import torch
  from grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  from torch import nn
  from torch.optim import AdamW
  from torch.utils.data import BatchSampler, DataLoader, RandomSampler
  from torch.utils.data.distributed import DistributedSampler
  from tqdm import tqdm
  from transformers import get_linear_schedule_with_warmup
  from utils import logger
label_desc: Use image node_id 4 for calling the Distiller obj's underlying member method code with example usage: obj.round_batch(x, lengths) and returns: x,  lengths, x,  lengths