uvos/QRotaryTraining
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

working full training

Browse source
This commit is contained in:
uvos 2024-03-09 10:03:37 +01:00
parent 7a47fcdcc0
commit 11ea9eeaa7
2 changed files with 25 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

5
convertinglinear.py
View file

@ -2,11 +2,12 @@ import torch
class ConvertingLinear(torch.nn.Linear):
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None, output_dtype=None):
super().__init__(in_features, out_features, bias, device, dtype)
self.output_dtype = output_dtype
def forward(self, input: torch.Tensor):
output_dtype = input.dtype
output_dtype = input.dtype if self.output_dtype is None else self.output_dtype
output_device = input.device
if input.device != self.weight.device:
input = input.to(self.weight.device)