I want to get the network weights as float16, so the network is currently nn().to(torch.float16)
nd, however, I get an error due to the biases being float32 type.
This is the error
return F.conv2d(input, weight, bias, self.stride,
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Input type (c10::Half) and bias type (float) should be the same
this is the simplified net:
# %%
class NeuralNetwork(nn.Module):
"""
A PyTorch neural network for image classification.
Args:
channels: Number of channels in the input image.
height: Height of the input image in pixels.
width: Width of the input image in pixels.
"""
def __init__(self, channels: int, height: int, width: int):
super().__init__()
self.backbone = Backbone()
output_dimensions = self.backbone.forward(
torch.zeros(1, channels, height, width).to(device, dtype=dtype)
).shape
self.head = Head(output_dimensions[-1])
def forward(self, x):
x = self.backbone(x)
x = self.head(x)
return x
Does this need explicitly passing the dtype to each layer ? Is there a better way to achieve this ?