Repository navigation
It is recommended to remove the "device" parameter from the code #70
Description
Activity
In addition, the tensors defined in the class will not be actively transferred to "device", and they should be changed to "nn. Parameter", the content comes from AI, please fully verify before adopting
example:
origin:self.a = torch.zeros(1, dtype=torch.float32, requires_grad=True) self.b = torch.zeros(1, dtype=torch.float32, requires_grad=True) self.alpha = torch.zeros(1, dtype=torch.float32, requires_grad=True)change
self.a = nn.Parameter(torch.zeros(1, dtype=torch.float32)) self.b = nn.Parameter(torch.zeros(1, dtype=torch.float32)) self.alpha = nn.Parameter(torch.zeros(1, dtype=torch.float32))Thank you for the suggestion! We will look into this issue and improve the library.
The lightning framework supports the use of multiple devices, and users who do not use the framework can use the to method to move to the desired device, and the "device" parameter makes the lightning framework use RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0!
So I suggest removing the "device" parameter