Write timers to a tensorboard writer

If distributed is initialized, print only on rank 0.

Return arguments.

Delete existing index and create a new

training one epoch

Copied from torch.optim.optimizer

Start the timer.

Stop the timer.

Set lr for all parameters groups.

Return tokenizer.

Return timers.

Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.

save the metrics to the logging dictionary arguments: log_dict (Dict) -- metrics/media to be logged in current step

Customized load.

Loop through the registered actions and fire all callbacks Args: hook The name of the hook to check, defaults to all

Log a group of timers.

Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model

Allocate a chunk of memory from the buffer to tensor and copy the values.

Set random seed for reproducability.

Creates a specified YOLOv5 model Arguments: name (str): name of model, i.e. 'yolov5s' pretrained (bool): load pretrained weights

Init method based on N(0, sigma).

Reduce a tensor of losses across all GPUs.

Return my rank for the tensor model parallel group.

Return world size for the tensor model parallel group.

Helper method to get the model chunk ID given the iteration number.

Get the tensor model parallel group the caller rank belongs to.

Simple scaling.

Init method based on N(0, sigma/sqrt(2*num_layers).

Communicate tensors between stages. Used as helper method in other communication methods that are used in megatron/schedules.py. Takes the fo

Simple linear layer with weight initialization.