# .numel() returns the total number of elements in a tensor
small_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
print(f"Small tensor shape: {small_tensor.shape}")
print(f"Small tensor elements: {small_tensor.numel()}")  # 2 * 3 = 6

matrix_2d = torch.randn(50, 20)
print(f"2D matrix elements: {matrix_2d.numel()}")  # 50 * 20 = 1000

# Our actual tensors
print(f"LHS elements: {lhs.numel()}")  # 128 * 512 = 65,536
print(f"RHS elements: {rhs.numel()}")  # 256 * 512 = 131,072