Class SimpleModel

tests/test_custom_offloading_utils.py:31–45

class SimpleModel(nn.Module):
    def __init__(self, num_blocks=16):
        super().__init__()
        self.blocks = nn.ModuleList([
            TransformerBlock(i)
        for i in range(num_blocks)])
    
    def forward(self, x):
        for block in self.blocks:
            x = block(x)
        return x

    @property
    def device(self):
        return next(self.parameters()).device


# Device Synchronization Tests