dont reshape non 4 dim arrays

This commit is contained in:
discountchubbs
2026-08-21 22:02:02 -07:00
parent df5695ba08
commit e372046ff1
@@ -186,14 +186,14 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
warped_dev = warped.to(Device.DEFAULT)
Tensor.realize(packed_npy_inputs_dev, warped_dev)
img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn).realize()
big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn).realize()
img = shift_and_sample(img_q, warped_dev[0:1], sample_skip_fn)
big_img = shift_and_sample(big_img_q, warped_dev[1:2], sample_skip_fn)
unpacked_tensors = [tensor.reshape(shape) for tensor, shape in zip(packed_npy_inputs_dev.split(npy_sizes), npy_shapes.values(), strict=True)]
unpacked_dict = dict(zip(npy_shapes.keys(), unpacked_tensors, strict=True))
desire_dev = unpacked_dict['desire']
desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn).realize()
desire_buf = shift_and_sample(desire_q, desire_dev.reshape(1, 1, -1), sample_desire_fn)
inputs = {desire_key: desire_buf}
for key, tensor_val in unpacked_dict.items():
@@ -202,22 +202,22 @@ def make_run_policy(vision_runner, policy_runners: list, features_slice: slice,
if 'prev_feat' in unpacked_dict:
prev_feat_dev = unpacked_dict['prev_feat']
feat_buf = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn).realize()
inputs['features_buffer'] = feat_buf.reshape(input_shapes['features_buffer'])
feat_buf = shift_and_sample(feat_q, prev_feat_dev.reshape(1, 1, -1), sample_skip_fn)
inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
if vision_runner:
vision_out_cast = next(iter(vision_runner({road_key: img, wide_key: big_img}).values())).cast('float32').realize()
if 'features_buffer' not in inputs:
new_feat = vision_out_cast[:, features_slice].reshape(1, -1).unsqueeze(0)
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn).realize()
inputs['features_buffer'] = feat_buf.reshape(input_shapes['features_buffer'])
inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
policy_outs = [next(iter(pol_runner(inputs).values())).cast('float32').realize() for pol_runner in policy_runners]
return (vision_out_cast, *policy_outs) if len(policy_outs) > 1 else (vision_out_cast, policy_outs[0])
inputs.update({road_key: img, wide_key: big_img})
if 'features_buffer' not in inputs:
feat_buf = sample_skip_fn(feat_q)
inputs['features_buffer'] = feat_buf.reshape(input_shapes['features_buffer'])
inputs['features_buffer'] = feat_buf if len(fb := input_shapes['features_buffer']) <= 3 else feat_buf.reshape(fb)
policy_out = next(iter(policy_runners[0](inputs).values())).cast('float32').realize()
if 'features_buffer' not in inputs and features_slice is not None: