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 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 grouped_batch_sampler import GroupedBatchSampler, create_lengths_groups
  from lm_seqs_dataset import LmSeqsDataset
  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