It may not be as simple as this, but when I look at the layers documentation I see this for MSELossFlat:
class FlattenedLoss():
"Same as `func`, but flattens input and target."
def __init__(self, func, *args, axis:int=-1, floatify:bool=False, is_2d:bool=True, **kwargs):
self.func,self.axis,self.floatify,self.is_2d = func(*args,**kwargs),axis,floatify,is_2d
def __repr__(self): return f"FlattenedLoss of {self.func}"
@property
def reduction(self): return self.func.reduction
@reduction.setter
def reduction(self, v): self.func.reduction = v
def __call__(self, input:Tensor, target:Tensor, **kwargs)->Rank0Tensor:
input = input.transpose(self.axis,-1).contiguous()
target = target.transpose(self.axis,-1).contiguous()
if self.floatify: target = target.float()
input = input.view(-1,input.shape[-1]) if self.is_2d else input.view(-1)
return self.func.__call__(input, target.view(-1), **kwargs)
def MSELossFlat(*args, axis:int=-1, floatify:bool=True, **kwargs):
"Same as `nn.MSELoss`, but flattens input and target."
return FlattenedLoss(nn.MSELoss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)
If I wanted to use MAE loss instead would it be as simple as this?
def MAELossFlat(*args, axis:int=-1, floatify:bool=True, **kwargs):
“Same as nn.L1Loss
, but flattens input and target.”
return FlattenedLoss(nn.L1Loss, *args, axis=axis, floatify=floatify, is_2d=False, **kwargs)
and if it is that simple, where do I make that change to add it? From looking at the layers.py file it appears to be part of the source code.