Tensorboard callback not working

I’m running into the error Boolean value of Tensor with no values is ambiguous while using TensorBoardCallback. This was happening even with a fresh install of the packages using pip (first encountered this while using the editable version).

Here’s the error

RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_13950/1465911307.py in <module>
      1 cbs = [TensorBoardCallback(projector=True)]
----> 2 learn.fit_one_cycle(3, cbs=cbs)

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/callback/schedule.py in fit_one_cycle(self, n_epoch, lr_max, div, div_final, pct_start, wd, moms, cbs, reset_opt)
    111     scheds = {'lr': combined_cos(pct_start, lr_max/div, lr_max, lr_max/div_final),
    112               'mom': combined_cos(pct_start, *(self.moms if moms is None else moms))}
--> 113     self.fit(n_epoch, cbs=ParamScheduler(scheds)+L(cbs), reset_opt=reset_opt, wd=wd)
    114 
    115 # Cell

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/learner.py in fit(self, n_epoch, lr, wd, cbs, reset_opt)
    219             self.opt.set_hypers(lr=self.lr if lr is None else lr)
    220             self.n_epoch = n_epoch
--> 221             self._with_events(self._do_fit, 'fit', CancelFitException, self._end_cleanup)
    222 
    223     def _end_cleanup(self): self.dl,self.xb,self.yb,self.pred,self.loss = None,(None,),(None,),None,None

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/learner.py in _with_events(self, f, event_type, ex, final)
    161 
    162     def _with_events(self, f, event_type, ex, final=noop):
--> 163         try: self(f'before_{event_type}');  f()
    164         except ex: self(f'after_cancel_{event_type}')
    165         self(f'after_{event_type}');  final()

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/learner.py in __call__(self, event_name)
    139 
    140     def ordered_cbs(self, event): return [cb for cb in self.cbs.sorted('order') if hasattr(cb, event)]
--> 141     def __call__(self, event_name): L(event_name).map(self._call_one)
    142 
    143     def _call_one(self, event_name):

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastcore/foundation.py in map(self, f, gen, *args, **kwargs)
    152     def range(cls, a, b=None, step=None): return cls(range_of(a, b=b, step=step))
    153 
--> 154     def map(self, f, *args, gen=False, **kwargs): return self._new(map_ex(self, f, *args, gen=gen, **kwargs))
    155     def argwhere(self, f, negate=False, **kwargs): return self._new(argwhere(self, f, negate, **kwargs))
    156     def argfirst(self, f, negate=False): return first(i for i,o in self.enumerate() if f(o))

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastcore/basics.py in map_ex(iterable, f, gen, *args, **kwargs)
    664     res = map(g, iterable)
    665     if gen: return res
--> 666     return list(res)
    667 
    668 # Cell

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastcore/basics.py in __call__(self, *args, **kwargs)
    649             if isinstance(v,_Arg): kwargs[k] = args.pop(v.i)
    650         fargs = [args[x.i] if isinstance(x, _Arg) else x for x in self.pargs] + args[self.maxi+1:]
--> 651         return self.func(*fargs, **kwargs)
    652 
    653 # Cell

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/learner.py in _call_one(self, event_name)
    143     def _call_one(self, event_name):
    144         if not hasattr(event, event_name): raise Exception(f'missing {event_name}')
--> 145         for cb in self.cbs.sorted('order'): cb(event_name)
    146 
    147     def _bn_bias_state(self, with_bias): return norm_bias_params(self.model, with_bias).map(self.opt.state)

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/callback/core.py in __call__(self, event_name)
     43                (self.run_valid and not getattr(self, 'training', False)))
     44         res = None
---> 45         if self.run and _run: res = getattr(self, event_name, noop)()
     46         if event_name=='after_fit': self.run=True #Reset self.run to True at each end of fit
     47         return res

~/miniconda3/envs/tt/lib/python3.9/site-packages/fastai/callback/tensorboard.py in before_fit(self)
     57             b = self.dls.one_batch()
     58             self.learn._split(b)
---> 59             self.writer.add_graph(self.model, *self.xb)
     60 
     61     def after_batch(self):

~/miniconda3/envs/tt/lib/python3.9/site-packages/torch/utils/tensorboard/writer.py in add_graph(self, model, input_to_model, verbose)
    725         if hasattr(model, 'forward'):
    726             # A valid PyTorch model should have a 'forward' method
--> 727             self._get_file_writer().add_graph(graph(model, input_to_model, verbose))
    728         else:
    729             # Caffe2 models do not have the 'forward' method

~/miniconda3/envs/tt/lib/python3.9/site-packages/torch/utils/tensorboard/_pytorch_graph.py in graph(model, args, verbose)
    292             raise e
    293 
--> 294     if verbose:
    295         print(graph)
    296     list_of_nodes = parse(graph, trace, args)

RuntimeError: Boolean value of Tensor with no values is ambiguous

This problem could also be on the pytorch side, and I noticed recent fixes in their source code for this file, but not sure if they were solving this exact problem.