CustomDataParallel

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

Bases: DataParallel

Wrapper class for multi-gpu training.

from https://github.com/pytorch/tutorials/issues/836