initializing c,h with with dimension [1,out_features] in case of batch_size 1 inside forward function, but not #ones(mode_mask)

[shravansingh123]

We are using Dataloader without collate_fn, which mean with one tree at a time. Inside forward function, the batch_size will be 1(becuase of node_order.shape[0]) and that will be used to initialize c,h with tensor of [1,out_features]. But inside _run_lstm function we are doing c[node_mask, :] = i * u , which will throw File "D:\phd\ijcai\treelstm\tree_lstm.py", line 127, in _run_lstm print(c[node_mask, :].shape)

IndexError: too many indices for tensor of dimension 2

Because we are passing node_mask which could have many 1's depending upon how many nodes will be operated at iteration 0. i*u is of size [NumberOfOnes(node_mask),out_features].

Please suggest.