Out of memory issue with large channel tensors

[nada-ak]

I managed to obtain the channel's frequency response h_freq which has a shape of TensorShape([1, 126, 1, 9, 128, 14, 300]) corresponding to [num_steps, num_ut, num_ut_ant, num_bs, num_bs_ant, num_ofdm_symbols, num_subcarriers]. This is a very large tensor to handle on my local hardware. I am currently facing issues with applying regularized zero-forcing precoding to my channel as such

precoded_channel = RZFPrecodedChannel(resource_grid=resource_grid, stream_management=stream_management)
h_eff = precoded_channel(h_freq, tx_power=dl_power, alpha=no) 

For completeness, dl_power has a shape of TensorShape([1, 9, 14, 14, 600]) which corresponds to [batch_size, num_bs, num_streams_per_bs, num_ofdm_sym, num_subcarriers].

I am wondering if there is a safe way to batch h_freq into chunks for easier processing through the precoded channel. I run on CPU and my RAM size is 16 GB.