CustomDataParallel¶

class behavenet.models.base.CustomDataParallel(module, device_ids=None, output_device=None, dim=0)[source]¶

Bases: torch.nn.parallel.data_parallel.DataParallel

Wrapper class for multi-gpu training.