src/nn/embedding.rs:92:
pub struct PositionalEncoding { table: Tensor, dim: usize }
...
self.table.to(input.device()).index_select(&positions)
table is a bare Tensor, not a Buffer, so named_buffers() is empty and
Module::to_device never moves it. It sits on the host forever and every forward
copies max_len * dim floats host to device just to slice out the seq_len rows
it actually wants. At max_len = 4096, dim = 512 that's 8 MB per call.
Put the table in a Buffer and register it in named_buffers(), or cache the
device copy.
While in there, forward never asserts input.last_dim() == self.dim.
src/nn/embedding.rs:92:tableis a bareTensor, not aBuffer, sonamed_buffers()is empty andModule::to_devicenever moves it. It sits on the host forever and every forwardcopies
max_len * dimfloats host to device just to slice out theseq_lenrowsit actually wants. At
max_len = 4096, dim = 512that's 8 MB per call.Put the table in a
Bufferand register it innamed_buffers(), or cache thedevice copy.
While in there,
forwardnever assertsinput.last_dim() == self.dim.