Not sure if I’m doing right. Am I doing it right for the sample code below?
def main():
    model = function(1000)
    print(torch.cuda.device_count())
    if True:
        if torch.cuda.device_count() > 1:
            model.model = MyDataParallel(model.model)

    start = time.monotonic()
    s = torch.cuda.current_stream()
    e_start = torch.cuda.Event(enable_timing=True)
    e_finish = torch.cuda.Event(enable_timing=True)
    s.record_event(e_start)

    model.train()

    torch.cuda.synchronize()
    s.record_event(e_finish)
    e_finish.synchronize()
    end = time.monotonic()

    print('Forward latency is: {} '.format(e_start.elapsed_time(e_finish)))
    print("end - start = ", end - start)

if __name__ == '__main__':
    main()