def _record_function_with_param(self): u = torch.randn(3, 4, 5, requires_grad=True) with _profile(with_stack=True, use_kineto=kineto_available(), record_shapes=True) as prof: with record_function("## TEST 1 ##", "1, 2, 3"): rf_handle = _record_function_with_args_enter("## TEST 2 ##", 1, False, 2.5, [u, u], "hello", u) _record_function_with_args_exit(rf_handle) return prof
def payload(self, use_cuda=False): u = torch.randn(3, 4, 5, requires_grad=True) with record_function("## TEST 1 ##", "1, 2, 3"): rf_handle = _record_function_with_args_enter("## TEST 2 ##", 1, False, 2.5, [u, u], (u, u), "hello", u) x = torch.randn(10, 10, requires_grad=True) if use_cuda: x = x.cuda() y = torch.randn(10, 10, requires_grad=True) if use_cuda: y = y.cuda() z = x + y + x * y + x * y z.backward(z) if use_cuda: z = z.cpu() _record_function_with_args_exit(rf_handle)