if hidden is None:
return {
"memory": cuda(T.zeros(b, m, w).fill_(δ), gpu_id=self.gpu_id),
"link_matrix": cuda(T.zeros(b, 1, m, m), gpu_id=self.gpu_id),
"precedence": cuda(T.zeros(b, 1, m), gpu_id=self.gpu_id),
"read_weights": cuda(T.zeros(b, r, m).fill_(δ), gpu_id=self.gpu_id),
After Change
"usage_vector": cuda(T.zeros(b, m), gpu_id=self.gpu_id)
}
// Build FLANN randomized k-d tree indexes for each batch
hx = rebuild_indexes(hx)
else:
// hidden["memory"] = hidden["memory"].clone()
hidden["link_matrix"] = hidden["link_matrix"].clone()
hidden["precedence"] = hidden["precedence"].clone()