Hi

first of all, thank you for your code

however,

Q1. shouldn't 'out_dim = featuers.shape[0]' be modified to 'out_dim = features.shape[1]' in the split_stack function?

Q2. shouldn't 'our[tar_idx] = features' be modified to 'out[tar_idx] += features' in the stack function?