added flattenedLoss func

app.py

@@ -26,6 +26,30 @@ import gradio as gr
 from fastai.vision.all import *
 import skimage
 
+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
+        functools.update_wrapper(self, self.func)
+
+    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()
+            
+        # Label smoothing experiment
+        target = (target * 0.9 + 0.05)
+        target[:,0] = 1
+
+        input = input.view(-1,input.shape[-1]) if self.is_2d else input.view(-1)
+        return self.func.__call__(input, target.view(-1), **kwargs)
+
 learn = load_learner('stage2-resnet50-size256-export.pkl')
 
 labels = learn.dls.vocab