logic fix for result allocate things
diff --git a/aten/src/aten/function_wrapper.py b/aten/src/aten/function_wrapper.py
index 9f914a4..5b2fcc8 100644
--- a/aten/src/aten/function_wrapper.py
+++ b/aten/src/aten/function_wrapper.py
@@ -208,15 +208,15 @@
def emit_body(env,option):
body = []
- for arg in option['formals_list']:
- if requires_checked_cast(arg):
- check_cast = CHECKED_CAST[arg['type']].substitute(env,arg_name=arg['name'])
- body.append("auto {}_ = {};".format(arg['name'],check_cast))
for arg in option['arguments']:
- if arg.get('allocate',False):
- body.append(
- CodeTemplate('auto ${arg_name}_ = new ${Tensor}(context);').substitute(
- env,arg_name=arg['name']))
+ if requires_checked_cast(arg):
+ if arg.get('allocate',False):
+ body.append(
+ CodeTemplate('auto ${arg_name}_ = new ${Tensor}(context);').substitute(
+ env,arg_name=arg['name']))
+ else:
+ check_cast = CHECKED_CAST[arg['type']].substitute(env,arg_name=arg['name'])
+ body.append("auto {}_ = {};".format(arg['name'],check_cast))
option['actuals'] = processor_type_env['state'] + get_arguments(option)
call = CodeTemplate("${THTensor}_${cname}(${actuals})").substitute(env)